LlamaIndex

pip install llama-index                 # llama-index-core + the OpenAI LLM and embedding integrations
pip install llama-index-readers-file    # PDF, DOCX, PPTX, CSV parsing for SimpleDirectoryReader
pip install llama-index-vector-stores-chroma llama-index-retrievers-bm25 llama-index-postprocessor-cohere-rerank

from llama_index.core import Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# Both read OPENAI_API_KEY from the environment. Left unset, the defaults are the
# older gpt-3.5-turbo and text-embedding-ada-002.
Settings.llm = OpenAI(model="gpt-4o-mini")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# Local models instead: pip install llama-index-llms-ollama llama-index-embeddings-ollama
# from llama_index.llms.ollama import Ollama
# from llama_index.embeddings.ollama import OllamaEmbedding
# Settings.llm = Ollama(model="llama3.1", request_timeout=360.0, context_window=8000)
# Settings.embed_model = OllamaEmbedding(model_name="nomic-embed-text")

from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex, load_index_from_storage

documents = SimpleDirectoryReader(
    "docs/handbook", recursive=True, required_exts=[".md", ".pdf"], filename_as_id=True
).load_data()
index = VectorStoreIndex.from_documents(documents)   # chunks, embeds, stores in memory
index.storage_context.persist(persist_dir="storage")

# Later runs: reload instead of re-embedding
index = load_index_from_storage(StorageContext.from_defaults(persist_dir="storage"))
response = index.as_query_engine(similarity_top_k=4).query("How many vacation days do new hires get?")
print(response)
for hit in response.source_nodes:
    print(hit.score, hit.metadata["file_name"])

import os
import chromadb
from llama_index.core.extractors import TitleExtractor
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.vector_stores.chroma import ChromaVectorStore

collection = chromadb.PersistentClient(path="chroma_db").get_or_create_collection("handbook")
vector_store = ChromaVectorStore(chroma_collection=collection)

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=1024, chunk_overlap=200),
        TitleExtractor(),            # per document: one LLM call per chunk (first 5) plus one to merge; adds document_title
        Settings.embed_model,        # must be a stage when a vector store is attached
    ],
    docstore=SimpleDocumentStore(),  # remembers a hash per document id
    vector_store=vector_store,
)
if os.path.exists("pipeline_storage"):
    pipeline.load("pipeline_storage")
nodes = pipeline.run(documents=documents)   # unchanged files are skipped, edited ones are replaced
pipeline.persist("pipeline_storage")
index = VectorStoreIndex.from_vector_store(vector_store)

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.postprocessor.cohere_rerank import CohereRerank
from llama_index.retrievers.bm25 import BM25Retriever

chunks = vector_store.get_nodes(node_ids=None)   # every chunk in the Chroma collection
hybrid = QueryFusionRetriever(
    [index.as_retriever(similarity_top_k=10), BM25Retriever.from_defaults(nodes=chunks, similarity_top_k=10)],
    similarity_top_k=10,
    num_queries=1,               # 1 = no LLM-generated query variations
    mode="reciprocal_rerank",
)
engine = RetrieverQueryEngine.from_args(
    hybrid,
    node_postprocessors=[CohereRerank(top_n=4)],   # reads COHERE_API_KEY
)

import asyncio
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.tools import QueryEngineTool
from llama_index.core.workflow import Context

handbook_tool = QueryEngineTool.from_defaults(
    query_engine=engine,
    name="handbook_search",
    description="Answers questions about the employee handbook: time off, expenses, hardware.",
)
agent = FunctionAgent(
    tools=[handbook_tool],
    llm=Settings.llm,
    system_prompt="Answer HR questions. Search the handbook before answering.",
)

async def main():
    ctx = Context(agent)   # carries chat history between runs
    print(await agent.run("How many vacation days do new hires get?", ctx=ctx))
    print(await agent.run("And how many of them carry over?", ctx=ctx))

asyncio.run(main())

from llama_index.core.evaluation import BatchEvalRunner, FaithfulnessEvaluator, RelevancyEvaluator

async def evaluate(engine, questions):
    judge = OpenAI(model="gpt-4o")
    runner = BatchEvalRunner(
        {"faithfulness": FaithfulnessEvaluator(llm=judge), "relevancy": RelevancyEvaluator(llm=judge)},
        workers=4,
    )
    results = await runner.aevaluate_queries(engine, queries=questions)
    for name, evals in results.items():
        print(name, sum(e.passing for e in evals) / len(evals))

import sys
from llama_index.core.query_engine import CitationQueryEngine

engine = CitationQueryEngine.from_args(index, similarity_top_k=4, citation_chunk_size=512)
response = engine.query(sys.argv[1])
print(response)
for i, hit in enumerate(response.source_nodes, start=1):
    print(f"[{i}] {hit.metadata['file_name']}")

python ask_handbook.py "How many vacation days do new hires get, and do unused days carry over?"

New hires receive 25 vacation days per calendar year, prorated from the start date [1].
Up to 5 unused days carry over and expire on March 31 [1].
[1] time-off.md
[2] expenses.md

from datetime import date
from llama_index.core.agent.workflow import FunctionAgent, ToolCallResult
from llama_index.core.tools import QueryEngineTool

def days_between(start: str, end: str) -> int:
    """Number of days between two ISO dates (YYYY-MM-DD)."""
    return (date.fromisoformat(end) - date.fromisoformat(start)).days

tools = [
    QueryEngineTool.from_defaults(handbook_index.as_query_engine(), name="handbook_search",
        description="HR policies: time off, expenses, hardware purchases."),
    QueryEngineTool.from_defaults(runbook_index.as_query_engine(), name="runbook_search",
        description="Engineering runbooks: deploys, incident response, on-call."),
    days_between,
]
agent = FunctionAgent(tools=tools, llm=Settings.llm, system_prompt="Use the tools; cite which one you used.")

async def ask(question: str):
    handler = agent.run(question)
    async for event in handler.stream_events():
        if isinstance(event, ToolCallResult):
            print("tool:", event.tool_name, event.tool_kwargs)
    print(await handler)