aimu.rag

Retrieval-augmented generation primitives: plain functions over the MemoryStore interface. See the RAG how-to.

Splitting

aimu.rag.split_text

split_text(text: str, *, chunk_size: int = 1000, chunk_overlap: int = 200, separators: list[str] | None = None, length_function: Callable[[str], int] = len) -> list[str]

Split text into overlapping chunks.

Parameters:

Name	Type	Description	Default
text	str	The text to split.	required
chunk_size	int	Maximum size of each chunk, measured by length_function.	1000
chunk_overlap	int	How much trailing content of one chunk repeats at the start of the next (continuity across boundaries). Must be smaller than chunk_size.	200
separators	list[str] | None	Separator hierarchy, largest semantic unit first. Defaults to paragraphs -> lines -> sentences -> words -> characters.	None
length_function	Callable[[str], int]	Measures chunk size. Defaults to len (characters); pass e.g. a tokenizer's encode-and-count for token-aware chunking.	len

Returns:

Type	Description
list[str]	A list of chunk strings in document order. Empty / whitespace-only input returns
list[str]	[].

Pipeline

aimu.rag.ingest

ingest(store: Any, documents: Union[str, list[str]], *, chunk_size: int = 1000, chunk_overlap: int = 200, separators: list[str] | None = None, length_function: Callable[[str], int] = len) -> int

Split documents into chunks and store each chunk in store.

chunk_size, chunk_overlap, separators, and length_function are forwarded to :func:~aimu.rag.split_text.

Parameters:

Name	Type	Description	Default
store	Any	Any object with a store(content: str) method (a :class:~aimu.memory.MemoryStore).	required
documents	Union[str, list[str]]	A single document string or a list of them.	required

Returns:

Type	Description
int	The number of chunks stored.

aimu.rag.retrieve

retrieve(store: Any, query: str, *, n_results: int = 5, **search_kwargs: Any) -> list[str]

Return up to n_results stored chunks relevant to query.

A thin, RAG-named pass-through to store.search(). Extra keyword arguments are forwarded to the store (e.g. max_distance= for :class:~aimu.memory.SemanticMemoryStore). To rerank, fetch a wider set here and pass it through :func:~aimu.rag.rerank.

aimu.rag.format_context

format_context(chunks: list[str], *, separator: str = '\n\n', numbered: bool = False) -> str

Join retrieved chunks into a single context string for prompt augmentation.

With numbered=True each chunk is prefixed [1], [2] ... (useful when you want the model to cite which chunk it used).

Reranking

aimu.rag.rerank

Optional cross-encoder reranking for retrieved chunks.

A retriever (vector or keyword search) is fast but coarse. Reranking re-scores each candidate against the query with a cross-encoder -- slower but more accurate -- so you fetch a wide set, then keep the best few. Backed by sentence-transformers (the [hf] extra). Plain function, not a class.

rerank

rerank(query: str, documents: list[str], *, model: str = DEFAULT_RERANK_MODEL, top_n: Optional[int] = None) -> list[str]

Reorder documents by cross-encoder relevance to query, most relevant first.

Parameters:

Name	Type	Description	Default
query	str	The search query.	required
documents	list[str]	Candidate chunks (e.g. the output of :func:~aimu.rag.retrieve with a generous n_results).	required
model	str	A sentence-transformers CrossEncoder model id.	DEFAULT_RERANK_MODEL
top_n	Optional[int]	Keep only the top N after reranking (default: keep all, reordered).	None

Returns:

Type	Description
list[str]	The documents reordered by descending relevance, truncated to top_n if set.
list[str]	An empty input returns [] without loading the model.