Skip to content

LlamaIndex

LlamaIndex is a data framework for LLM applications focused on document ingestion, indexing, and retrieval. It excels at building RAG systems over heterogeneous research data — PDFs, notebooks, code, databases.

Installation

pip install llama-index llama-index-llms-openai-like llama-index-embeddings-openai python-dotenv
# Optional vector stores:
pip install llama-index-vector-stores-chroma chromadb
pip install llama-index-vector-stores-faiss faiss-cpu

Setup

import os
from dotenv import load_dotenv
from llama_index.core import Settings
from llama_index.llms.openai_like import OpenAILike
from llama_index.embeddings.openai import OpenAIEmbedding

load_dotenv()

Settings.llm = OpenAILike(
    model="llama-3-1-8b",
    api_base=os.environ["PAIS_API_BASE"],
    api_key=os.environ["PAIS_API_KEY"],
    is_chat_model=True,
    context_window=128000,
    max_tokens=2048,
)

Settings.embed_model = OpenAIEmbedding(
    model="qwen3-vl-embedding-8b",
    api_base=os.environ["PAIS_API_BASE"],
    api_key=os.environ["PAIS_API_KEY"],
)

Once Settings is configured, all LlamaIndex components use PAIS automatically.

Index a document collection

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

# Load PDFs, markdown, text files from a directory
documents = SimpleDirectoryReader("./papers/").load_data()

# Build index (embeds all documents using PAIS embedding model)
index = VectorStoreIndex.from_documents(documents, show_progress=True)

# Persist to disk
index.storage_context.persist(persist_dir="./index_storage")

Load a persisted index:

from llama_index.core import StorageContext, load_index_from_storage

storage_context = StorageContext.from_defaults(persist_dir="./index_storage")
index = load_index_from_storage(storage_context)

Query the index

query_engine = index.as_query_engine(similarity_top_k=4)

response = query_engine.query(
    "What evaluation metrics are used across these papers and how are they compared?"
)
print(response)
print("\nSource nodes:")
for node in response.source_nodes:
    print(f"  [{node.score:.3f}] {node.metadata.get('file_name', 'unknown')}")

Chat engine (multi-turn Q&A)

chat_engine = index.as_chat_engine(
    chat_mode="context",
    verbose=True,
    system_prompt=(
        "You are a research assistant with access to a collection of ML papers. "
        "Cite specific papers when answering."
    ),
)

response = chat_engine.chat("What approaches are used for few-shot learning?")
print(response)
response = chat_engine.chat("Which of those works best for low-resource NLP tasks?")
print(response)

Sub-question query engine

For complex questions that span multiple documents, the sub-question engine breaks the query into sub-queries:

from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.tools import QueryEngineTool

query_tool = QueryEngineTool.from_defaults(
    query_engine=index.as_query_engine(),
    name="research_papers",
    description="A collection of ML research papers on model training and evaluation",
)

sub_question_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=[query_tool],
    use_async=True,
)

response = sub_question_engine.query(
    "Compare the datasets used and the model sizes across all the papers."
)
print(response)

Chroma persistent vector store

For large corpora, use Chroma as a persistent vector store instead of in-memory:

import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext

chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("research_papers")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

Agentic tool use

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
import arxiv

def search_arxiv(query: str, max_results: int = 5) -> str:
    """Search arXiv for recent papers. Returns titles and abstracts."""
    client = arxiv.Client()
    search = arxiv.Search(query=query, max_results=max_results,
                          sort_by=arxiv.SortCriterion.SubmittedDate)
    results = []
    for paper in client.results(search):
        results.append(f"Title: {paper.title}\nAbstract: {paper.summary[:400]}...")
    return "\n\n".join(results)

arxiv_tool = FunctionTool.from_defaults(fn=search_arxiv)
index_tool = QueryEngineTool.from_defaults(
    query_engine=index.as_query_engine(),
    name="local_papers",
    description="Search the local paper collection",
)

agent = ReActAgent.from_tools(
    [arxiv_tool, index_tool],
    verbose=True,
    max_iterations=10,
)

response = agent.chat(
    "Find recent arXiv papers on mixture-of-experts, then compare them to "
    "what's in our local collection."
)
print(response)

See also