tai-aitutor — Migrating from LlamaIndex

Plain-Python building blocks for RAG: provider-neutral LLM calls, embeddings, chunking, retrieval, and evaluation. Small flat modules, plain functions, readable source. If you can't read a module in one sitting, that's a bug.

Coming from LlamaIndex? What you know carries over; the surface is what changes. Where you set Settings, built an Index, handed a question to a QueryEngine and scored it with an Evaluator, you now call configure(), open a Chroma collection, run search() then build_rag_prompt() + generate(), and score with hit_rate / reciprocal_rank / the judge functions. Each step you knew as a class is a function here. The migration map below pairs them symbol by symbol — read it as a translation table, not a rebuild.

The package supplies the parts inside each step. It does not supply the pipeline: the indexing loop and the retrieve → prompt → generate sequence stay in your code, where you can see and change them. Composites that would hide several steps behind one call are deliberately absent.

Install

Requires Python 3.12+ (Colab's current runtime, so the floor stays there while it does).

# the three providers
pip install "tai-aitutor[gemini,openai,anthropic]"

# + the vector store
pip install "tai-aitutor[gemini,openai,anthropic,rag]"

cohere (reranking, Cohere embeddings) and sentence-transformers (local embeddings) are optional heavies: the package never imports them unless you call rerank / embed_cohere / embed_local. Install them only if you use those — via the extras ([rerank], [local]) or directly (pip install cohere sentence-transformers); any recent version works.

Quickstart

from tai_aitutor import (
    build_rag_prompt, chunk_document, configure, embed, generate,
    get_collection, load_csv, search, setup_notebook, show_answer,
)

IN_COLAB = setup_notebook(required_keys=("GOOGLE_API_KEY",))   # Colab Secrets or .env
PROVIDER = "gemini"   # @param ["gemini", "openai", "anthropic"]
configure(provider=PROVIDER)                                    # replaces Settings

docs = load_csv("articles.csv", text_col="content",
                meta_cols=("title", "url", "source"), id_col="title")
col = get_collection("kb", path="./db")                         # the collection IS the index

chunks = [c for d in docs for c in chunk_document(d)]           # chunk → embed → upsert:
col.add(                                                        # three steps, all visible
    ids=[c.id for c in chunks],
    documents=[c.text for c in chunks],
    embeddings=embed([c.text for c in chunks]),
    metadatas=[c.metadata for c in chunks],
)

question = "What is RAG?"
hits = search(question, col, top_k=5)                           # retrieve
show_answer(generate(build_rag_prompt(question, hits)), hits)   # prompt → generate

Loading data

The package ships no dataset URLs and no downloaders. Download however you like, then hand the file to load_csv:

docs = load_csv("articles.csv", text_col="content",
                meta_cols=("title", "url", "source"), id_col="title")

Other formats are four lines of your own code, and reading them is worth more than a wrapper. QADataset.load(path) opens LlamaIndex's EmbeddingQAFinetuneDataset JSON byte-for-byte, so existing eval files work unchanged.

Migration map: LlamaIndex → tai_aitutor

Delete on sight: every nest_asyncio.apply() call (it existed only for LlamaIndex's async), every llama-index-* pip pin, and LLAMA_CLOUD_API_KEY setup.

Where a row says the replacement is "your own code", that is the point — the operation was small enough that a wrapper cost more than it saved.

Config

LLMs and embeddings

Documents, chunking, ingestion, storage

Retrieval and answering

Evaluation

Agents, chat, tools, routing

Fine-tuning

No package equivalent — fine-tuning is a training script, not a building block.

What's in the package

config (configure, setup_notebook, require_keys, in_colab) · llm (generate, extract) · embeddings (embed, embed_cohere, embed_local, EMBED_DIM) · tokens (n_tokens) · documents (Document, load_csv) · chunking (Chunk, chunk, chunk_document, chunk_sentences, heading_aware_markdown_chunks, sentence_window_chunks) · vectorstore (get_collection, reset_collection, get_all_chunks, build_where_filter) · retrieval (search, ScoredChunk, expand_window, code_tokenize, BM25Index, rrf_fuse, rerank, judge_rerank, rewrite_query, hyde_search, decompose_question) · synthesis (build_rag_prompt) · evals (QADataset, hit_rate, reciprocal_rank, evaluate_retrieval, sweep_top_k, the three judges) · tools (Tool, tool, search_web) · router (route, RouteDecision) · display · errors (all subclass ValueError)

Runnable end-to-end examples in examples/.

Design rules

1. Functions over object graphs. State lives in one configure() call and in the objects you pass explicitly — the collection, the index, the dataset. Where LlamaIndex handed you an object carrying the pipeline inside it, you hold the pieces.
2. No hidden composites. If a call would bundle several steps you should be able to see, it isn't here.
3. No data downloads in the package. It loads the file you hand it, nothing more.
4. Loud failures. Typo'd kwargs raise TypeError, data parsing never uses eval(), missing keys and missing extras say exactly what to install or set.

Development

pip install -e ".[gemini,openai,anthropic,rag,rerank]" pytest ruff
pytest              # fully offline — provider SDKs are faked
ruff check src tests

Releases: bump __version__ in src/tai_aitutor/__init__.py, commit, then git tag vX.Y.Z && git push --tags — .github/workflows/release.yml publishes to PyPI via Trusted Publishing.

License

MIT — see LICENSE.