RAG Module

RAG pipeline API reference

​

RAG Module

The RAG module provides a thin orchestration layer over Knowledge for retrieval-augmented generation.

​

Core Classes

​

RAG

The main pipeline class that orchestrates retrieval and generation.

from praisonaiagents.rag import RAG, RAGConfig

rag = RAG(
    knowledge=knowledge,           # Required: Knowledge instance
    llm=None,                      # Optional: LLM instance
    config=RAGConfig(),            # Optional: Configuration
    reranker=None,                 # Optional: Reranker instance
    context_builder=None,          # Optional: Custom context builder
    citation_formatter=None,       # Optional: Custom citation formatter
)

​

Methods

query(question, **kwargs) -> RAGResult

Execute a RAG query and return result with citations.

result = rag.query("What is the main finding?")
print(result.answer)
print(result.citations)

aquery(question, **kwargs) -> RAGResult

Async version of query.

result = await rag.aquery("What is the conclusion?")

aquery runs retrieval and generation on a worker thread via asyncio.to_thread, so the event loop stays responsive and contextvars set on the awaiting task are visible inside the thread. aretrieve does the same for retrieval only. See Async Safety.

stream(question, **kwargs) -> Iterator[str]

Stream response tokens.

for chunk in rag.stream("Summarize the document"):
    print(chunk, end="", flush=True)

astream(question, **kwargs) -> AsyncIterator[str]

Async streaming.

async for chunk in rag.astream("Explain the methodology"):
    print(chunk, end="", flush=True)

Retrieval and reranking are offloaded via asyncio.to_thread. When the LLM is sync-only (no astream method), the fallback generator is pumped from a worker thread and reuses the retrieved context — retrieval runs exactly once per astream call.

retrieve(query, top_k=None, rerank=None, **kwargs) -> ContextPack

Retrieve context without generating an answer. top_k and rerank override the config for this call only.

from praisonaiagents.rag import RAG, RAGConfig

rag = RAG(knowledge=knowledge, config=RAGConfig(top_k=5, rerank=True))

# One-off override — does not change rag.config
pack = rag.retrieve("niche query", top_k=1, rerank=False)
assert pack.metadata["top_k"] == 1          # effective value used
assert rag.config.top_k == 5                # config unchanged

Per-call overrides are scoped to the call. rag.retrieve(top_k=…, rerank=…) uses the override values and leaves rag.config unchanged, including under concurrent aretrieve calls. The effective values appear on ContextPack.metadata["top_k"] and ContextPack.metadata["reranked"].

get_citations(question, **kwargs) -> List[Citation]

Get citations without generating an answer.

citations = rag.get_citations("What sources mention X?")
for c in citations:
    print(f"[{c.id}] {c.source}: {c.text[:100]}")

​

RAGConfig

Configuration for the RAG pipeline. top_k and rerank are defaults — passing them to retrieve() overrides only that call and never rewrites RAGConfig.

from praisonaiagents.rag import RAGConfig, RetrievalStrategy

config = RAGConfig(
    top_k=5,                                    # Chunks to retrieve
    min_score=0.0,                              # Minimum relevance score
    max_context_tokens=4000,                    # Context token limit
    include_citations=True,                     # Include citations
    retrieval_strategy=RetrievalStrategy.BASIC, # Retrieval strategy
    rerank=False,                               # Enable reranking
    rerank_top_k=3,                             # Results after rerank
    template="...",                             # Prompt template
    system_prompt=None,                         # System prompt
    stream=False,                               # Stream by default
)

​

Retrieval Strategies

from praisonaiagents.rag import RetrievalStrategy

RetrievalStrategy.BASIC   # Simple vector search
RetrievalStrategy.FUSION  # Reciprocal rank fusion
RetrievalStrategy.HYBRID  # Dense + sparse retrieval

​

RAGResult

Result from a RAG query.

@dataclass
class RAGResult:
    answer: str                    # Generated answer
    citations: List[Citation]      # Source citations
    context_used: str              # Context passed to LLM
    query: str                     # Original query
    metadata: Dict[str, Any]       # Timing, stats, etc.

​

Properties

• has_citations - Boolean indicating if citations exist
• format_answer_with_citations() - Format answer with source references

​

Citation

Source citation for RAG answers.

@dataclass
class Citation:
    id: str                        # Citation ID (e.g., "1")
    source: str                    # Source document
    text: str                      # Text snippet
    score: float                   # Relevance score
    doc_id: Optional[str]          # Document identifier
    chunk_id: Optional[str]        # Chunk identifier
    offset: Optional[int]          # Character offset
    metadata: Dict[str, Any]       # Additional metadata

​

Protocols

The RAG module uses protocols for extensibility.

​

ContextBuilderProtocol

Custom context assembly logic.

from typing import Any, Dict, List, Union
from praisonaiagents.knowledge.models import SearchResultItem
from praisonaiagents.rag.protocols import ContextBuilderProtocol

ResultItem = Union[Dict[str, Any], SearchResultItem]

class MyContextBuilder:
    def build(
        self,
        results: List[ResultItem],
        max_tokens: int = 4000,
        deduplicate: bool = True,
    ) -> str:
        # Custom context building logic
        return assembled_context

​

CitationFormatterProtocol

Custom citation formatting.

from praisonaiagents.rag.protocols import CitationFormatterProtocol

class MyCitationFormatter:
    def format(
        self,
        results: List[Dict[str, Any]],
        start_id: int = 1,
    ) -> List[Citation]:
        # Custom citation formatting
        return citations

​

Context Utilities

Helper functions for context building that accept both dict and SearchResultItem formats.

Dict results

from praisonaiagents.rag import build_context, deduplicate_chunks

# Traditional dict format
results = [
    {"text": "First content", "metadata": {"filename": "a.pdf"}},
    {"text": "Second content", "metadata": {"filename": "b.pdf"}},
]

# Build context from dict results
context, used_results = build_context(
    results=results,
    max_tokens=4000,
    deduplicate=True,
    include_source=True,
)

# Remove duplicate chunks
unique = deduplicate_chunks(results)

SearchResultItem results

from praisonaiagents.rag import build_context, deduplicate_chunks
from praisonaiagents.knowledge.models import SearchResultItem

# SearchResultItem format
results = [
    SearchResultItem(text="First content", source="a.pdf", filename="a.pdf"),
    SearchResultItem(text="Second content", source="b.pdf", filename="b.pdf"),
]

# Mixed dict + object input also works
results.append({"text": "Third content", "metadata": {"filename": "c.pdf"}})

# Build context from SearchResultItem objects
context, used_results = build_context(
    results=results,
    max_tokens=2000,
    include_source=True,
)

unique = deduplicate_chunks(results)

build_context and deduplicate_chunks accept a mix of dict results and SearchResultItem objects. When include_source=True, the label is taken from metadata["filename"] / metadata["source"] first, falling back to the top-level filename / source attribute on the item, and finally to Source N.

​

Result Item Formats

The context utilities support two input formats with automatic fallback for metadata lookups:

Lookup	1st choice	2nd choice	Fallback
source	metadata["source"]	item.source (object) / item["source"] (dict)	""
filename	metadata["filename"]	item.filename / item["filename"]	""
text	item.text / item["text"]	item.memory / item["memory"]	"" (item skipped in build_context)

​

Integration with Knowledge

RAG uses Knowledge for all retrieval operations:

from praisonaiagents import Knowledge
from praisonaiagents.rag import RAG

# Knowledge handles indexing
knowledge = Knowledge()
knowledge.add("documents/")

# RAG handles answering
rag = RAG(knowledge=knowledge)
result = rag.query("What is discussed?")

​

Error Handling

from praisonaiagents.rag import RAG

rag = RAG(knowledge=knowledge)

try:
    result = rag.query("Question")
except Exception as e:
    print(f"RAG error: {e}")

At the retriever layer, failures are non-raising by default: a broken vector store returns empty chunks instead of throwing. Check RetrievalResult.retrieval_failed and .error to tell a real failure from a genuine no-match — see Detect Retrieval Failures.

​

Performance Tips

1. Batch indexing: Add multiple documents at once
2. Tune top_k: Start with 5, adjust based on quality
3. Use min_score: Filter low-relevance results
4. Enable reranking: For higher precision (costs latency)
5. Stream responses: Better UX for long answers