Implementing Vector Search with LlamaIndex

A comprehensive guide to using the LlamaIndex framework to implement vector search in LLM-based applications.

Table of Contents

1. Training Overview
2. Getting Started with the LlamaIndex Framework
   • What is LlamaIndex?
   • Use Cases
   • General Architecture
3. Installation and Configuration
   • Virtual Environment
   • Required Packages
   • OpenAI API Key
4. Module 1 — Using LlamaIndex as a Vector Store
   • Quickstart Example
   • Final Code
5. Module 2 — Creating a Chroma Index
   • Chroma Overview
   • Step 1: Create a Client and a Collection
   • Vector Embeddings Explained
   • Step 2: Define the Embedding Function
   • Step 3: Load Documents
   • Step 4: Configure the Chroma Vector Store
   • Step 5: Query the Data
   • Step 6: Customize Parameters — Cosine Similarity
   • Final Code
6. Module 3 — Implementing a Multi-step Query Pipeline
   • Workflows Overview
   • Basic Workflow: StartEvent and StopEvent
   • Workflow with Custom Events
   • The Complete Q&A Query Pipeline
   • Step by Step: Building the Pipeline
   • Complete Final Code
7. Summary and Best Practices
8. Quick Reference

1. Training Overview

This training covers three progressive major modules:

┌────────────────────────────────────────────────────────────┐
│           Implementing Vector Search with LlamaIndex        │
├────────────────────┬───────────────────┬───────────────────┤
│   Module 1 (13m)   │   Module 2 (23m)  │   Module 3 (43m)  │
│  LlamaIndex as a   │  Creating a       │  Multi-step       │
│   Vector Store     │  Chroma Index     │  Query Pipeline   │
└────────────────────┴───────────────────┴───────────────────┘

2. Getting Started with the LlamaIndex Framework

What is LlamaIndex?

LlamaIndex is an open source framework designed to simplify and optimize the connection between Large Language Models (LLMs) and your custom data. It acts as a data orchestration layer to enable vector search in your applications.

graph TB
    A[Custom Data<br/>Documents · Databases · APIs · Web] --> B[LlamaIndex<br/>Orchestration Layer]
    B --> C[LLMs<br/>OpenAI · Other models]
    C --> D[Intelligent Applications<br/>Context-aware · RAG]
    
    style B fill:#f4a261,stroke:#e76f51,color:#000
    style C fill:#457b9d,stroke:#1d3557,color:#fff
    style A fill:#a8dadc,stroke:#457b9d,color:#000
    style D fill:#2a9d8f,stroke:#264653,color:#fff

Problem solved by LlamaIndex:

LLMs are trained on large amounts of general data but do not know your specific data. LlamaIndex bridges this gap by giving LLMs access to your custom information: local documents, websites, databases, APIs.

Use Cases

mindmap
  root((LlamaIndex))
    Workflows
      Chain actions
      AI-driven automation
    AI Agents
      Autonomous search
      Decision making
    Chatbots
      Customer support
      Document Q&A
    RAG
      Context augmentation
      Augmented generation
    Query Engines
      Semantic search
      Multi-step

General Architecture

flowchart LR
    subgraph Sources["Data Sources"]
        D1[Documents]
        D2[Web pages]
        D3[Databases]
        D4[APIs]
    end

    subgraph LlamaIndex["LlamaIndex Framework"]
        L1[Data Connectors<br/>Loaders]
        L2[Index<br/>VectorStoreIndex]
        L3[Query Engine<br/>Retriever]
        L4[LLM Interface]
    end

    subgraph Output["Output"]
        O1[Augmented responses]
        O2[Summaries]
        O3[Generated content]
    end

    Sources --> L1
    L1 --> L2
    L2 --> L3
    L3 --> L4
    L4 --> Output

    style LlamaIndex fill:#fff3e0,stroke:#ff9800

LlamaIndex is:

• ✅ Open source and community-driven
• ✅ Modular and flexible
• ✅ Equipped with a simple high-level API
• ✅ Capable of processing custom data sources

3. Installation and Configuration

Virtual Environment

It is strongly recommended to create an isolated Python virtual environment for each project.

macOS

# Create the virtual environment
python3 -m venv venv

# Activate the virtual environment
source venv/bin/activate

# Deactivate
deactivate

Windows

# Create the virtual environment
python -m venv venv

# Activate the virtual environment
venv\Scripts\activate

# Deactivate
deactivate

Note: For Module 2, a specific Python version (3.10.6) may be required. Use pyenv to manage multiple Python versions.

# macOS — Install pyenv via Homebrew
brew install pyenv

# Install and use Python 3.10.6
pyenv install 3.10.6
pyenv global 3.10.6

# Verify the version
python -V

Required Packages

Module 1 — Quickstart

python-dotenv==1.1.0
colorama==0.4.6
llama-index==0.12.28
llama_index.readers.web

Installation:

# macOS
pip3 install -r requirements.txt
pip3 install llama_index.readers.web

# Windows
pip install -r requirements.txt
pip install llama_index.readers.web

Module 2 — Chroma Index

python-dotenv==1.1.0
colorama==0.4.6
chromadb
llama-index>=0.10.10
openai>=1.0.0
tqdm>=4.64.0
numpy<2

Installation:

# macOS
pip3 install -r requirements.txt
pip3 install llama-index-readers-web llama-index-vector-stores-chroma \
             llama-index-embeddings-langchain llama-index-embeddings-huggingface

# Windows
pip install -r requirements.txt
pip install llama-index-readers-web llama-index-vector-stores-chroma \
            llama-index-embeddings-langchain llama-index-embeddings-huggingface

Module 3 — Query Pipeline

python-dotenv==1.1.0
colorama==0.4.6
llama-index==0.12.28
llama_index.readers.web

Troubleshooting tip: If you get an OpenAI validation error, update the dependencies:

pip3 install --upgrade llama-index openai pydantic

OpenAI API Key

1. Create an account at platform.openai.com
2. Go to Profile → API Keys
3. Click Create new secret key and copy the key
4. Create a .env file at the project root:

# .env
OPENAI_API_KEY="your-secret-key"

5. Load the key in Python code:

from dotenv import load_dotenv
load_dotenv()

⚠️ The key can only be viewed once upon creation. Keep it somewhere safe. If it is compromised, revoke it and create a new one.

4. Module 1 — Using LlamaIndex as a Vector Store

Quickstart Example

The quickstart example shows how to:

1. Import the necessary classes
2. Load documents from a web page
3. Create a VectorStoreIndex
4. Create a query engine
5. Query the LLM with augmented context

sequenceDiagram
    participant U as User
    participant App as Python Application
    participant Web as Web Page<br/>(paulgraham.com)
    participant LI as LlamaIndex
    participant LLM as OpenAI LLM

    U->>App: Enters a question
    App->>Web: SimpleWebPageReader.load_data()
    Web-->>App: HTML documents → Text
    App->>LI: VectorStoreIndex.from_documents()
    LI->>LI: Converts to vector embeddings
    App->>LI: index.as_query_engine()
    App->>LI: query_engine.query(question)
    LI->>LLM: Sends context + question
    LLM-->>App: Augmented response
    App-->>U: Displays the response

Starter file (01_04/starter/main.py):

from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.readers.web import SimpleWebPageReader
from colorama import Fore, Style

load_dotenv()

def main():
    """
    Main function to load a webpage, build a vector-based query engine,
    and generate content with context
    """
    user_input = input(Fore.BLUE + "Q: " + Fore.RESET)

if __name__ == "__main__":
    main()

Final Code — Module 1

Final file (01_04/final/main.py):

from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex
from llama_index.readers.web import SimpleWebPageReader
from colorama import Fore, Style

load_dotenv()

def main():
    """
    Main function to load a webpage, build a vector-based query engine,
    and generate content with context
    """
    user_input = input(Fore.BLUE + "Q: " + Fore.RESET)

    # Load documents from a web page
    documents = SimpleWebPageReader(html_to_text=True).load_data(
        ["http://paulgraham.com/worked.html"]
    )

    # Build the vector index
    index = VectorStoreIndex.from_documents(documents)

    # Create the query engine
    query_engine = index.as_query_engine()

    # Generate a response with context
    response = query_engine.query(user_input)
    print(Fore.GREEN + "A: " + str(response))

if __name__ == "__main__":
    main()

To start the application:

# macOS
python3 main.py

# Windows
python main.py

Sample questions to ask:

1. What kind of things did Paul Graham build or create as a child?
2. How did his interest in art and computers evolve over time?
3. What was Paul’s experience like at RISD?
4. How did the acquisition of Viaweb by Yahoo impact his path?
5. Why did Paul Graham start Y Combinator?

5. Module 2 — Creating a Chroma Index

Chroma Overview

ChromaDB is an open source vector database that simplifies building LLM-powered applications by making knowledge and facts “pluggable” for language models.

flowchart TD
    A[Natural text<br/>e.g. Paul Graham essays] -->|SimpleWebPageReader| B[Documents]
    B -->|HuggingFaceEmbedding| C[Vector Embeddings\n0.23, -0.15, 0.87, ...]
    C -->|ChromaVectorStore| D[(ChromaDB\nVector Database)]
    E[User question] -->|Embedding| F[Query vector]
    F -->|Cosine similarity / L2| D
    D -->|Relevant documents| G[Augmented context]
    G -->|OpenAI LLM| H[Final response]

    style D fill:#9b59b6,stroke:#8e44ad,color:#fff
    style C fill:#3498db,stroke:#2980b9,color:#fff
    style H fill:#2ecc71,stroke:#27ae60,color:#000

Vector Embeddings Explained

A vector embedding is a numerical representation of textual content as a vector of floating-point numbers:

"LlamaIndex is an open source framework"
    ↓  (embedding model)
[0.23, -0.15, 0.87, 0.02, -0.44, 0.91, ...]

graph LR
    subgraph Vector space
        A((Text A))
        B((Text B))
        C((Text C))
    end
    
    A ---|Low distance<br/>= High similarity| B
    A ----|Large distance<br/>= Low similarity| C
    
    style A fill:#e74c3c,color:#fff
    style B fill:#e74c3c,color:#fff
    style C fill:#3498db,color:#fff

Step 1: Create a Client and a Collection

import chromadb

# Initialize an ephemeral Chroma client (in memory)
chroma_client = chromadb.EphemeralClient()

# Create a new collection named "search-index"
chroma_collection = chroma_client.create_collection("search-index")

An ephemeral client stores data in memory only (non-persistent). For persistence, use chromadb.PersistentClient(path="./chroma_db").

Step 2: Define the Embedding Function

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# Load the HuggingFace BAAI embedding model
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")

The BAAI/bge-base-en-v1.5 model is responsible for converting text into numerical vectors. This model enables the machine to understand and measure semantic similarity between texts.

Step 3: Load Documents

from llama_index.readers.web import SimpleWebPageReader

# Load documents from a web page
documents = SimpleWebPageReader(html_to_text=True).load_data(
    ["http://paulgraham.com/worked.html"]
)

SimpleWebPageReader is a Data Connector provided by LlamaIndex that reads documents directly from a web page. The html_to_text=True option converts HTML to plain text.

Step 4: Configure the Chroma Vector Store

from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext, VectorStoreIndex

# Configure the vector store with the Chroma collection
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# Create the storage context
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# Create the index from documents
# This step converts documents to embeddings and stores them in ChromaDB
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    embed_model=embed_model
)

💡 VectorStoreIndex.from_documents() automatically performs two operations:

1. Converts text to vector embeddings via the embed_model
2. Stores these vectors in ChromaDB via the vector_store

Step 5: Query the Data

from colorama import Fore

# Create the query engine
query_engine = index.as_query_engine()

# Generate an augmented response
response = query_engine.query(user_input)
print(Fore.GREEN + "A: " + str(response) + Fore.RESET)

Step 6: Customize Parameters — Cosine Similarity

What is Cosine Similarity?

Cosine similarity measures the cosine of the angle between two vectors in a vector space. It indicates how semantically similar two pieces of text are.

$$\text{cosine_similarity}(A, B) = \frac{A \cdot B}{|A| \cdot |B|}$$

graph LR
    subgraph "Distance Metrics"
        L2["L2 (Euclidean)\n— Chroma default\n— measures absolute distance"]
        COS["Cosine Similarity\n— angle between vectors\n— better for text"]
    end
    L2 -.->|"Replace with"| COS
    style COS fill:#2ecc71,stroke:#27ae60,color:#000
    style L2 fill:#e74c3c,stroke:#c0392b,color:#fff

Configure cosine similarity in ChromaDB

# Create the collection with the cosine metric
chroma_collection = chroma_client.create_collection(
    "search-index",
    metadata={"hnsw:space": "cosine"}
)

# Add a retriever with control over the number of results
retriever = index.as_retriever(similarity_top_k=3)
# similarity_top_k=3 → returns the 3 most relevant results

# Use the retriever to fetch relevant documents
nodes = retriever.retrieve(user_input)

By default, Chroma uses L2 (Euclidean) distance. For text, cosine similarity is generally better because it is insensitive to vector magnitude.

Final Code — Module 2

import chromadb
import warnings
import os
from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex
from llama_index.readers.web import SimpleWebPageReader
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from colorama import Fore

warnings.filterwarnings("ignore", category=UserWarning, module="chromadb")

load_dotenv()

def main():
    """Main function to run the ChromaDB vector search example."""

    user_input = input(Fore.BLUE + "Q: " + Fore.RESET)

    # Step 1: Create a client and a new collection
    chroma_client = chromadb.EphemeralClient()
    chroma_collection = chroma_client.create_collection("search-index")

    # Step 2: Define the embedding function
    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")

    # Step 3: Load documents from the web
    documents = SimpleWebPageReader(html_to_text=True).load_data(
        ["http://paulgraham.com/worked.html"]
    )

    # Step 4: Configure ChromaVectorStore and load the data
    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,
        embed_model=embed_model
    )

    # Step 5: Query the data and generate an augmented response
    query_engine = index.as_query_engine()
    response = query_engine.query(user_input)
    print(Fore.GREEN + "A: " + str(response) + Fore.RESET)

if __name__ == "__main__":
    main()

6. Module 3 — Implementing a Multi-step Query Pipeline

Workflows Overview

A workflow with LlamaIndex is an event-driven execution, divided into multiple steps, that chains actions to accomplish a complex task.

graph LR
    SE[StartEvent] --> E1[Event 1]
    E1 --> E2[Event 2]
    E2 --> E3[Event 3]
    E3 --> STE[StopEvent]
    
    style SE fill:#3498db,stroke:#2980b9,color:#fff
    style STE fill:#e74c3c,stroke:#c0392b,color:#fff
    style E1 fill:#f39c12,stroke:#e67e22,color:#fff
    style E2 fill:#9b59b6,stroke:#8e44ad,color:#fff
    style E3 fill:#2ecc71,stroke:#27ae60,color:#000

Key principles:

• Each step receives an event as input and returns an event as output
• The output of a step becomes the input of the next step
• The @step decorator indicates that a method belongs to the workflow
• Workflows run asynchronously with asyncio

Basic Workflow: StartEvent and StopEvent

StartEvent and StopEvent are two special events provided directly by LlamaIndex, ready to use.

File 03_02/final/main.py:

import asyncio
from dotenv import load_dotenv
from colorama import Fore
from llama_index.core.workflow import (
    Event,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)

load_dotenv()

class MyEvent(Event):
    response: str

class BasicWorkflow(Workflow):

    @step
    async def start(self, ev: StartEvent) -> MyEvent:
        print(Fore.BLUE + f"Starting workflow...{ev.topic}" + Fore.RESET)
        return MyEvent(response=str(ev.topic))

    @step
    async def stop(self, ev: MyEvent) -> StopEvent:
        print(Fore.GREEN + f"Stopping workflow...{ev.response}" + Fore.RESET)
        return StopEvent(result=str(ev.response))

async def main():
    w = BasicWorkflow(timeout=60, verbose=False)
    result = await w.run(topic="AI")
    print(result)

if __name__ == "__main__":
    asyncio.run(main())

sequenceDiagram
    participant M as main()
    participant W as BasicWorkflow
    participant SE as StartEvent
    participant ME as MyEvent
    participant STE as StopEvent

    M->>W: w.run(topic="AI")
    W->>SE: Created automatically
    SE->>W: start(ev: StartEvent)
    W->>ME: MyEvent(response="AI")
    ME->>W: stop(ev: MyEvent)
    W->>STE: StopEvent(result="AI")
    STE-->>M: "AI"

Workflow with Custom Events

Building a more advanced workflow that uses an OpenAI LLM to generate a joke with chained events.

File 03_03/final/main.py:

import asyncio
import os
from dotenv import load_dotenv
from colorama import Fore
from llama_index.llms.openai import OpenAI
from llama_index.core.workflow import (
    Event,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)

load_dotenv()

# Custom Events
class QueryEvent(Event):
    prompt: str

class GenerateEvent(Event):
    answer: str

class JokeWorkflow(Workflow):
    llm = OpenAI()

    @step
    async def query(self, ev: StartEvent) -> QueryEvent:
        prompt = f"Generate a joke about {ev.input}"
        print(Fore.BLUE + f"Asking AI to generate a joke: {ev.input}" + Fore.RESET)
        return QueryEvent(prompt=str(prompt))

    @step
    async def generate(self, ev: QueryEvent) -> GenerateEvent:
        response = await self.llm.acomplete(ev.prompt)
        print(Fore.GREEN + f"Generating a joke using AI: {response}" + Fore.RESET)
        return GenerateEvent(answer=str(response))

    @step
    async def stop(self, ev: GenerateEvent) -> StopEvent:
        print(Fore.GREEN + f"Stopping workflow...{ev.answer}" + Fore.RESET)
        return StopEvent(result=str(ev.answer))

async def main():
    w = JokeWorkflow(timeout=60, verbose=False)
    result = await w.run(input="bees")
    print(result)

if __name__ == "__main__":
    asyncio.run(main())

Event chain diagram:

stateDiagram-v2
    [*] --> StartEvent : w.run(input="bees")
    StartEvent --> QueryEvent : @step query()\nBuilds the prompt
    QueryEvent --> GenerateEvent : @step generate()\nLLM generates the joke
    GenerateEvent --> StopEvent : @step stop()\nReturns the result
    StopEvent --> [*] : Final result

The Complete Q&A Query Pipeline

The goal of Module 3 is to build a Q&A pipeline with vector search that:

1. Takes a question from a user
2. Generates optimized queries via an LLM
3. Retrieves relevant documents via vector search
4. Generates responses for each optimized query
5. Summarizes all responses into a single coherent final answer

flowchart TD
    A([🚀 StartEvent\nUser question]) --> B

    B["@step start()\nPasses the question"]
    B -->|QueryEvent| C

    C["@step expand_query()\nGenerates 3 optimized queries\nvia LLM"]
    C -->|MultiQueryEvent\noptimized_queries| D

    D["@step retrieve()\nVector Search\nsimilarity_top_k=3"]
    D -->|RetrievalEvent\nretrieved_docs| E

    E["@step generate()\nGenerates responses\nfor each document/query"]
    E -->|GenerateEvent\ngenerated_answers| F

    F["@step summarize()\nMerges all responses\ninto a single answer"]
    F -->|SummaryEvent| G

    G["@step end()\nReturns the final answer"]
    G --> H([🏁 StopEvent\nFinal answer])

    style A fill:#3498db,stroke:#2980b9,color:#fff
    style H fill:#e74c3c,stroke:#c0392b,color:#fff
    style C fill:#f39c12,stroke:#e67e22,color:#fff
    style D fill:#9b59b6,stroke:#8e44ad,color:#fff
    style E fill:#1abc9c,stroke:#16a085,color:#fff
    style F fill:#e67e22,stroke:#d35400,color:#fff

Step by Step: Building the Pipeline

Event Definitions

from typing import List
from llama_index.core.workflow import Event

# Event 1: Passes the initial question
class QueryEvent(Event):
    input: str

# Event 2: Contains the optimized queries generated by the LLM
class MultiQueryEvent(Event):
    optimized_queries: List[str]

# Event 3: Contains the documents retrieved by the retriever
class RetrievalEvent(Event):
    retrieved_docs: List[str]
    optimized_queries: List[str]

# Event 4: Contains the generated responses for each query/document
class GenerateEvent(Event):
    generated_answers: list  # List of dicts {"query": ..., "response": ...}

# Event 5: Contains the final summary
class SummaryEvent(Event):
    text: str

Step 1 — start: Initialization

@step
async def start(self, ev: StartEvent) -> QueryEvent:
    """Starts the workflow with the user's question."""
    return QueryEvent(input=str(ev.input))

Step 2 — expand_query: Query Optimization

@step
async def expand_query(self, queryEvent: QueryEvent) -> MultiQueryEvent:
    """Generates 3 optimized variants of the question via an LLM."""
    formatted_prompt = f"Generate 3 different ways to phrase the question: '{queryEvent.input}'"
    completion = await self.llm.acomplete(formatted_prompt)
    optimized_queries = [
        line.strip("- ")
        for line in completion.text.split("\n")
        if line.strip()
    ]
    return MultiQueryEvent(optimized_queries=optimized_queries)

Prompt Engineering: The prompt asks the LLM to generate 3 different formulations of the same question. This improves vector search coverage because different formulations may retrieve different relevant documents.

Step 3 — retrieve: Vector Search

@step
async def retrieve(self, ev: MultiQueryEvent) -> RetrievalEvent:
    """Retrieves relevant documents for each optimized query."""
    results = []
    for q in ev.optimized_queries:
        nodes = self.retriever.retrieve(q)       # Vector search
        results.extend([n.text for n in nodes])  # Extract text
    return RetrievalEvent(
        retrieved_docs=results,
        optimized_queries=ev.optimized_queries
    )

The retriever is built on the index: self.retriever = index.as_retriever(similarity_top_k=3). It uses vector embeddings to find the 3 semantically closest documents to each query.

Step 4 — generate: Response Generation

@step
async def generate(self, ev: RetrievalEvent) -> GenerateEvent:
    """Generates a response for each query/document combination."""
    prompt_template = PromptTemplate(
        "Generate a response for the given query:\n\n{query}\n\n"
        "Using the following context:\n\n{doc}"
    )

    summaries = []
    for query in ev.optimized_queries:
        for doc in ev.retrieved_docs:
            prompt = prompt_template.format(query=query, doc=doc)
            completion = await self.llm.acomplete(prompt)
            summaries.append({
                "query": query,
                "response": completion.text.strip()
            })
    return GenerateEvent(generated_answers=summaries)

Step 5 — summarize: Final Summary

@step
async def summarize(self, ev: GenerateEvent) -> SummaryEvent:
    """Merges all responses into a single coherent summary."""
    summary_texts = [entry["response"] for entry in ev.generated_answers]
    merged_prompt = PromptTemplate(
        "Combine the following summaries into one coherent summary:\n\n{summaries}"
    )
    combined_prompt = merged_prompt.format(summaries="\n\n".join(summary_texts))
    final_completion = await self.llm.acomplete(combined_prompt)
    return SummaryEvent(text=str(final_completion.text.strip()))

Step 6 — end: End of Workflow

@step
async def end(self, ev: SummaryEvent) -> StopEvent:
    """Ends the workflow and returns the final answer."""
    return StopEvent(result=str(ev.text))

Complete Final Code — Module 3

File 03_09/final/main.py — complete version with colored outputs:

import asyncio
import os
from colorama import Fore, Style
from dotenv import load_dotenv
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.llms.openai import OpenAI
from llama_index.core.prompts import PromptTemplate
from llama_index.readers.web import SimpleWebPageReader
from llama_index.core import VectorStoreIndex
from llama_index.core.workflow import (
    Event,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)
from typing import List

load_dotenv()

# ─────────────────────────────────────────
# 1. Event Definitions
# ─────────────────────────────────────────

class QueryEvent(Event):
    input: str
    def __str__(self):
        return (f"{Fore.CYAN}==========❓QueryEvent ========== {Fore.RESET}\n"
                f"Starting Workflow with a query: {self.input}...")

class MultiQueryEvent(Event):
    optimized_queries: List[str]
    def __str__(self):
        return (f"{Fore.MAGENTA}==========❓MultiQueryEvent (OPTIMIZED QUERIES) ========== "
                f"{Fore.RESET}\nExpanded Queries: {', '.join(self.optimized_queries)}")

class RetrievalEvent(Event):
    retrieved_docs: List[str]
    optimized_queries: List[str]
    def __str__(self):
        return (f"{Fore.BLUE}========== 🔍 RetrievalEvent (RETRIEVED DOCUMENTS) ========== "
                f"{Fore.RESET}\nRetrieved {len(self.retrieved_docs)} documents.")

class GenerateEvent(Event):
    generated_answers: list  # List of dicts {"query": ..., "response": ...}
    def __str__(self):
        answers_str = "\n".join(
            f"Query: {entry['query']}\nResponse: {entry['response']}"
            for entry in self.generated_answers
        )
        return (f"{Fore.CYAN}========== 🖊️ GenerateEvent (ANSWERS) ========== "
                f"{Fore.RESET}\n\nAnswers:\n{answers_str}")

class SummaryEvent(Event):
    text: str
    def __str__(self):
        return (f"{Fore.GREEN}========== 🔤 SummaryEvent (SUMMARIES) ========== "
                f"{Fore.RESET}\n\nSummary: {self.text}{Fore.RESET}")

# ─────────────────────────────────────────
# 2. Workflow Definition
# ─────────────────────────────────────────

class QAWorkflow(Workflow):
    """
    Q&A Workflow with vector search capabilities.
    Optimizes queries, retrieves documents, and generates a summarized answer.
    """
    def __init__(self, index, timeout=None, max_retries=None, verbose=False):
        super().__init__()
        self.llm = OpenAI(model="gpt-3.5-turbo")
        self._timeout = timeout
        self._max_retries = max_retries
        self.index = index
        self.retriever = index.as_retriever(similarity_top_k=3)
        self.summarizer = TreeSummarize(llm=self.llm)

    @step
    async def start(self, ev: StartEvent) -> QueryEvent:
        """Starts the workflow with the initial question."""
        return QueryEvent(input=str(ev.input))

    @step
    async def expand_query(self, queryEvent: QueryEvent) -> MultiQueryEvent:
        """Generates 3 optimized formulations of the question."""
        print(queryEvent)
        formatted_prompt = (
            f"Generate 3 different ways to phrase the question: '{queryEvent.input}'"
        )
        completion = await self.llm.acomplete(formatted_prompt)
        optimized_queries = [
            line.strip("- ")
            for line in completion.text.split("\n")
            if line.strip()
        ]
        return MultiQueryEvent(optimized_queries=optimized_queries)

    @step
    async def retrieve(self, ev: MultiQueryEvent) -> RetrievalEvent:
        """Performs vector search for each optimized query."""
        print(ev)
        results = []
        for q in ev.optimized_queries:
            nodes = self.retriever.retrieve(q)
            results.extend([n.text for n in nodes])
        return RetrievalEvent(
            retrieved_docs=results,
            optimized_queries=ev.optimized_queries
        )

    @step
    async def generate(self, ev: RetrievalEvent) -> GenerateEvent:
        """Generates contextualized responses for each query/document pair."""
        print(ev)
        prompt_template = PromptTemplate(
            "Generate a response for the given query:\n\n{query}\n\n"
            "Using the following context:\n\n{doc}"
        )
        summaries = []
        for query in ev.optimized_queries:
            for doc in ev.retrieved_docs:
                prompt = prompt_template.format(query=query, doc=doc)
                completion = await self.llm.acomplete(prompt)
                summaries.append({
                    "query": query,
                    "response": completion.text.strip()
                })
        return GenerateEvent(generated_answers=summaries)

    @step
    async def summarize(self, ev: GenerateEvent) -> SummaryEvent:
        """Merges all responses into a coherent final summary."""
        print(ev)
        summary_texts = [entry["response"] for entry in ev.generated_answers]
        merged_prompt = PromptTemplate(
            "Combine the following summaries into one coherent summary:\n\n{summaries}"
        )
        combined_prompt = merged_prompt.format(
            summaries="\n\n".join(summary_texts)
        )
        final_completion = await self.llm.acomplete(combined_prompt)
        return SummaryEvent(text=str(final_completion.text.strip()))

    @step
    async def end(self, ev: SummaryEvent) -> StopEvent:
        """Ends the workflow and returns the final answer."""
        print(ev)
        return StopEvent(result=str(ev.text))

# ─────────────────────────────────────────
# 3. Main Entry Point
# ─────────────────────────────────────────

async def main():
    user_input = input("Enter your question: ")

    # Load data from the web
    documents = SimpleWebPageReader(html_to_text=True).load_data(
        ["https://paulgraham.com/greatwork.html"]
    )

    # Build the vector index
    index = VectorStoreIndex.from_documents(documents)

    # Initialize and run the workflow
    w = QAWorkflow(index=index, timeout=60, max_retries=3, verbose=False)
    print("========== START OF THE WORKFLOW ==========\n\n")
    result = await w.run(input=user_input)
    print("========== END OF THE WORKFLOW ==========")
    print("\n" + Fore.GREEN + "🧠 Answer:", result)

if __name__ == "__main__":
    asyncio.run(main())

Sample console output:

========== START OF THE WORKFLOW ==========

==========❓QueryEvent ==========
Starting Workflow with a query: what did the author do as a kid...

==========❓MultiQueryEvent (OPTIMIZED QUERIES) ==========
Expanded Queries: How did the author spend their childhood,
                  What activities did the author engage in during their youth?,
                  Can you describe the author's childhood experiences?

========== 🔍 RetrievalEvent (RETRIEVED DOCUMENTS) ==========
Retrieved 9 documents.

========== 🖊️ GenerateEvent (ANSWERS) ==========
[Generated responses for each query/document combination...]

========== 🔤 SummaryEvent (SUMMARIES) ==========
Summary: [Final consolidated summary...]

========== END OF THE WORKFLOW ==========

🧠 Answer: [Concise and coherent final answer]

7. Summary and Best Practices

Comparison of Approaches

graph TD
    A[Goal] --> B{Project complexity?}
    B -->|Simple / Prototype| C[Module 1\nSimple VectorStoreIndex\nQuickstart]
    B -->|Persistence required| D[Module 2\nChromaDB Index\nHuggingFace Embeddings]
    B -->|Advanced pipeline| E[Module 3\nMulti-step Workflow\nFull RAG]

    C --> C1["✓ Fast to implement\n✓ Ideal for testing\n✗ No persistence"]
    D --> D1["✓ Data persistence\n✓ Optimized embeddings\n✓ Cosine similarity"]
    E --> E1["✓ Full RAG\n✓ Optimized queries\n✓ Automatic summaries\n✗ More complex"]

    style C fill:#3498db,color:#fff
    style D fill:#9b59b6,color:#fff
    style E fill:#e74c3c,color:#fff

Summary Table of Key Classes

Best Practices

1. Always use a virtual environment to isolate dependencies
2. Never commit your OpenAI API key — use .env + .gitignore
3. Use cosine similarity rather than L2 for text processing tasks
4. similarity_top_k: adjust as needed (3-5 is generally optimal)
5. Asynchronous: LlamaIndex workflows use asyncio — remember await
6. Prompt Engineering: well-crafted prompts = better LLM responses
7. Manage versions: ChromaDB and LlamaIndex may require specific versions (e.g., Python 3.10.6 for ChromaDB)

RAG (Retrieval-Augmented Generation) Concept

flowchart LR
    Q[User\nquestion] --> E1[Question\nembedding]
    E1 --> VS[(Vector Store\nChromaDB)]
    VS -->|Top-K relevant\ndocuments| CTX[Augmented\ncontext]
    CTX --> P[Prompt\n= Question + Context]
    P --> LLM[LLM\nOpenAI]
    LLM --> R[Contextualized\nresponse]

    DOC[Source documents] --> E2[Document\nembedding]
    E2 --> VS

    style VS fill:#9b59b6,color:#fff
    style LLM fill:#e74c3c,color:#fff
    style R fill:#2ecc71,color:#000

8. Quick Reference

Essential Commands

# Create the virtual environment
python -m venv venv               # Windows
python3 -m venv venv              # macOS

# Activate
venv\Scripts\activate             # Windows
source venv/bin/activate          # macOS

# Install dependencies
pip install -r requirements.txt

# Run an example
python main.py                    # Windows
python3 main.py                   # macOS

# Update dependencies (if issues arise)
pip install --upgrade llama-index openai pydantic

Recommended Project Structure

my-project/
├── .env                    ← API key (do not commit)
├── .env.example            ← .env template (commit this)
├── .gitignore              ← Include .env and venv/
├── requirements.txt        ← Dependencies
├── main.py                 ← Entry point
└── venv/                   ← Virtual environment (do not commit)

Basic LlamaIndex Pattern

from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex
from llama_index.readers.web import SimpleWebPageReader

load_dotenv()

# 1. Load documents
documents = SimpleWebPageReader(html_to_text=True).load_data(["URL"])

# 2. Build the index
index = VectorStoreIndex.from_documents(documents)

# 3. Query
query_engine = index.as_query_engine()
response = query_engine.query("Your question here")
print(response)

Basic Workflow Pattern

import asyncio
from llama_index.core.workflow import Event, StartEvent, StopEvent, Workflow, step

class MyEvent(Event):
    data: str

class MyWorkflow(Workflow):
    @step
    async def step1(self, ev: StartEvent) -> MyEvent:
        return MyEvent(data=ev.input)

    @step
    async def step2(self, ev: MyEvent) -> StopEvent:
        return StopEvent(result=ev.data)

async def main():
    w = MyWorkflow(timeout=60)
    result = await w.run(input="Hello")
    print(result)

asyncio.run(main())

Search Terms

vector · search · llamaindex · rag · embeddings · artificial · intelligence · generative · ai · query · chroma · pipeline · workflow · cosine · similarity · configure · generation · index · pattern · quickstart · store