Ask questions about your own documents with an agent

What you get

A small Python agent that reads the files in a data folder, indexes them in a VectorStoreIndex (OpenAI embeddings), and answers questions about them. It also has a multiply tool, so one question can mix document search and maths. This is the official LlamaIndex starter tutorial with OpenAI.

Prerequisites

Python 3.10 or newer (llama-index 0.14.25 needs >=3.10, <4.0)
An OpenAI API key in OPENAI_API_KEY
A data folder with the documents you want to ask about (the tutorial uses a Paul Graham essay)

Steps

Install LlamaIndex. The llama-index package is a starter bundle that includes llama-index-core and the OpenAI LLM and embedding integrations.

Shell

pip install llama-index==0.14.25

Set your OpenAI key.

Shell

# MacOS/Linux
export OPENAI_API_KEY=XXXXX
# Windows
set OPENAI_API_KEY=XXXXX

Put some documents in a data folder. To use the tutorial's sample essay:

Shell

mkdir data
wget https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt -O data/paul_graham_essay.txt

Your folder should look like this:

Text

├── starter.py
└── data
    └── paul_graham_essay.txt

Create starter.py:

Python

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI
import asyncio
import os

# Create a RAG tool using LlamaIndex
documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

def multiply(a: float, b: float) -> float:
    """Useful for multiplying two numbers."""
    return a * b

async def search_documents(query: str) -> str:
    """Useful for answering natural language questions about an personal essay written by Paul Graham."""
    response = await query_engine.aquery(query)
    return str(response)

# Create an enhanced workflow with both tools
agent = FunctionAgent(
    tools=[multiply, search_documents],
    llm=OpenAI(model="gpt-4o-mini"),
    system_prompt="""You are a helpful assistant that can perform calculations
    and search through documents to answer questions.""",
)

# Now we can ask questions about the documents or do calculations
async def main():
    response = await agent.run(
        "What did the author do in college? Also, what's 7 * 8?"
    )
    print(response)

# Run the agent
if __name__ == "__main__":
    asyncio.run(main())

If you use your own files, change the search_documents docstring so it describes them. The agent reads that docstring to decide when to search.

Test it

Shell

python starter.py

The agent should answer what the author did in college (from the essay, through search_documents) and that 7 * 8 is 56 (through multiply).

Not confirmed

The docs show the expected output only for the multiply-only version, not for this RAG version, so the exact wording of the answer is not confirmed.
With no embed model set, VectorStoreIndex uses OpenAI's text-embedding-ada-002. The docs recommend passing OpenAIEmbedding(model="text-embedding-3-small") for new projects, but the starter code does not show where to put it, so it is left at the default here.
The download step uses wget, as the docs do. The docs give no alternative for machines without wget (you can save the file into data any other way).
The starter code imports os but does not use it; it is kept as written in the docs.