From raw OpenAI prompts to LlamaIndex query engines

Moving from raw OpenAI prompts to LlamaIndex query engines isn’t just a refactor — it’s a shift from prompt engineering to system engineering. The openai prompts to llamaindex query engine migration pays off when you need retrieval, structured output, streaming, and observability without building all that yourself. This guide walks through the migration end to end with runnable code at each step.

Step 1: Inventory your current OpenAI usage

Before adding dependencies, catalog what your raw calls actually do. Most teams find three patterns:

# Pattern A: Simple chat completion
response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": prompt}],
    temperature=0.2,
)

# Pattern B: Chat with few-shot examples baked into messages
messages = [
    {"role": "system", "content": "You are a SQL expert..."},
    {"role": "user", "content": "Schema: ...\nQuestion: ..."},
    {"role": "assistant", "content": "SELECT ..."},
    {"role": "user", "content": "Schema: ...\nQuestion: ..."},
]
response = client.chat.completions.create(model="gpt-4o-mini", messages=messages)

# Pattern C: Function calling for structured output
tools = [{"type": "function", "function": {"name": "extract", "parameters": schema}}]
response = client.chat.completions.create(
    model="gpt-4o-mini", messages=messages, tools=tools, tool_choice="auto"
)

Verify: Grep your codebase for chat.completions.create and completions.create. Note which calls use tools, streaming, or custom parameters like logprobs. You’ll map each pattern to a LlamaIndex equivalent.

Step 2: Install the minimal dependency set

LlamaIndex is modular. Install only what you need:

# Core + OpenAI integration + vector store (pick one)
pip install llama-index llama-index-llms-openai llama-index-embeddings-openai

# If you need a vector store (most migrations do)
pip install llama-index-vector-stores-chroma  # or pinecone, weaviate, qdrant, etc.

# Optional: observability, evaluation, structured output
pip install llama-index-callbacks-arize-phoenix  # or langfuse, wandb
pip install llama-index-output-parser-pydantic

Verify: Run python -c "import llama_index; print(llama_index.__version__)" and confirm no import errors.

Step 3: Replace the raw client with an LlamaIndex LLM wrapper

The OpenAI class wraps the same HTTP calls but adds retry logic, token counting, and callback hooks.

from llama_index.llms.openai import OpenAI
from llama_index.core.llms import ChatMessage, MessageRole

llm = OpenAI(
    model="gpt-4o-mini",
    temperature=0.2,
    max_tokens=1024,
    # These map 1:1 to OpenAI SDK params
    api_key="sk-...",  # or rely on OPENAI_API_KEY env var
    # LlamaIndex-specific: enable callbacks for token accounting
    callback_manager=None,  # inject later for observability
)

# Pattern A equivalent
messages = [ChatMessage(role=MessageRole.USER, content=prompt)]
response = llm.chat(messages)
print(response.message.content)

# Pattern B equivalent (few-shot)
messages = [
    ChatMessage(role=MessageRole.SYSTEM, content="You are a SQL expert..."),
    ChatMessage(role=MessageRole.USER, content="Schema: ...\nQuestion: ..."),
    ChatMessage(role=MessageRole.ASSISTANT, content="SELECT ..."),
    ChatMessage(role=MessageRole.USER, content="Schema: ...\nQuestion: ..."),
]
response = llm.chat(messages)

Verify: Run a test prompt. Confirm response.message.content matches your old output. Check response.raw for the underlying OpenAI response object if you need logprobs or finish_reason.

Step 4: Add retrieval — the real reason to migrate

Raw prompts stuff context into the message. LlamaIndex separates retrieval from synthesis. This is where the openai prompts to llamaindex query engine migration pays off.

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

# Configure global settings (or pass per-component)
Settings.llm = llm
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# Load documents — replace with your data source
documents = SimpleDirectoryReader("./data").load_data()

# Option A: In-memory (dev only)
index = VectorStoreIndex.from_documents(documents)

# Option B: Persistent Chroma (production)
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("docs")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
index = VectorStoreIndex.from_documents(documents, vector_store=vector_store)

# Create a query engine — this replaces your prompt + context stuffing
query_engine = index.as_query_engine(
    similarity_top_k=4,
    response_mode="compact",  # "tree_summarize" for multi-doc synthesis
    streaming=False,
)

response = query_engine.query("What does the API return on 429?")
print(response.response)  # synthesized answer
print(response.source_nodes)  # retrieved chunks with scores

Verify: Query something answerable only from your documents. Confirm source_nodes returns relevant chunks with similarity scores. If results look wrong, adjust similarity_top_k or check embedding model matches your query language.

Step 5: Enable streaming for latency-sensitive UIs

Streaming in LlamaIndex uses async generators. Wrap the query engine, not the LLM directly.

from llama_index.core.query_engine import StreamingAgentChatResponse

# Enable streaming on the query engine
streaming_engine = index.as_query_engine(
    similarity_top_k=4,
    response_mode="compact",
    streaming=True,
)

# Sync streaming (blocks but yields tokens)
streaming_response = streaming_engine.query("Explain the retry logic")
for token in streaming_response.response_gen:
    print(token, end="", flush=True)
print()

# Async streaming (for FastAPI, etc.)
async def stream_answer(question: str):
    response = await streaming_engine.aquery(question)
    async for token in response.async_response_gen():
        yield token

# FastAPI endpoint example
# @app.get("/ask")
# async def ask(q: str):
#     return StreamingResponse(stream_answer(q), media_type="text/plain")

Verify: Hit the endpoint or run the sync loop. Tokens should appear incrementally. Confirm source_nodes is still populated after streaming completes (streaming_response.source_nodes).

Step 6: Structured output with Pydantic output parsers

Replace function calling with LlamaIndex’s output parsers — they handle retries and validation automatically.

from pydantic import BaseModel, Field
from llama_index.core.output_parsers import PydanticOutputParser
from llama_index.core.query_engine import CitationQueryEngine

class APIErrorResponse(BaseModel):
    error_code: str = Field(description="HTTP status code as string")
    message: str = Field(description="Human-readable error message")
    retry_after_seconds: int | None = Field(default=None, description="Seconds to wait before retry")

parser = PydanticOutputParser(output_cls=APIErrorResponse)

# Use a query engine that enforces the schema
structured_engine = index.as_query_engine(
    similarity_top_k=3,
    response_mode="compact",
    output_parser=parser,
)

# The prompt automatically includes format instructions
response = structured_engine.query("What happens on rate limit?")
parsed: APIErrorResponse = response.response  # already validated
print(parsed.error_code, parsed.message, parsed.retry_after_seconds)

Verify: Feed a query that should trigger the schema. Confirm parsed is a real APIErrorResponse instance, not a string. Check response.metadata for parser retry count if validation failed initially.

Step 7: Add observability — tokens, latency, retrieval quality

LlamaIndex’s callback system captures everything. Hook it up once.

from llama_index.core import Settings, set_global_handler
from llama_index.callbacks import CallbackManager, TokenCountingHandler
import tiktoken

# Token counting (local, no external dependency)
token_counter = TokenCountingHandler(
    tokenizer=tiktoken.encoding_for_model("gpt-4o-mini").encode
)

# Or use an external platform (Phoenix, Langfuse, etc.)
# from llama_index.callbacks import ArizePhoenixCallbackHandler
# phoenix = ArizePhoenixCallbackHandler(endpoint="http://localhost:6006")

Settings.callback_manager = CallbackManager([token_counter])

# Run a query
response = query_engine.query("What's the rate limit policy?")

# Inspect counts
print(f"Prompt tokens: {token_counter.prompt_llm_token_count}")
print(f"Completion tokens: {token_counter.completion_llm_token_count}")
print(f"Embedding tokens: {token_counter.total_embedding_token_count}")

# Reset for next request
token_counter.reset_counts()

Verify: Run a known query. Token counts should be non-zero and match expectations (embedding tokens ≈ doc count × chunk size). If using Phoenix/Langfuse, open the UI and confirm traces show retrieval + synthesis spans.

Step 8: Evaluation — know when retrieval fails

Migration isn’t done until you can measure quality. LlamaIndex includes evaluators for faithfulness and relevancy.

from llama_index.core.evaluation import FaithfulnessEvaluator, RelevancyEvaluator

faithfulness = FaithfulnessEvaluator(llm=llm)
relevancy = RelevancyEvaluator(llm=llm)

# Evaluate a single response
eval_result = faithfulness.evaluate_response(response=response)
print(f"Faithful: {eval_result.passing}, Score: {eval_result.score}")
print(f"Feedback: {eval_result.feedback}")

# Batch evaluation over a test set
test_questions = [
    "What is the rate limit?",
    "How do I authenticate?",
    "What regions are supported?",
]

for q in test_questions:
    resp = query_engine.query(q)
    faith = faithfulness.evaluate_response(resp)
    rel = relevancy.evaluate_response(resp, query=q)
    print(f"Q: {q[:50]}... Faithful: {faith.passing} Relevant: {rel.passing}")

Verify: Run on 10-20 representative questions. Faithfulness < 0.7 usually means retrieval is pulling noise or the synthesizer is hallucinating. Relevancy < 0.7 means the query isn’t matching the right chunks — tune similarity_top_k, chunk size, or embedding model.

Step 9: Production hardening — timeouts, fallbacks, caching

Raw OpenAI calls need the same hardening. LlamaIndex gives you extension points.

import httpx
from llama_index.llms.openai import OpenAI
from llama_index.core.query_engine import RetryQueryEngine

# Custom HTTP client with timeouts
http_client = httpx.Client(timeout=httpx.Timeout(connect=5.0, read=30.0, write=10.0))

llm = OpenAI(
    model="gpt-4o-mini",
    temperature=0.2,
    http_client=http_client,
    # LlamaIndex retries on 429/5xx by default; customize:
    max_retries=3,
    retry_jitter=True,
)

# Wrap query engine with automatic retry on failure
retry_engine = RetryQueryEngine(
    query_engine=query_engine,
    max_retries=2,
    # Optional: different engine for retry (e.g., smaller model)
    # retry_query_engine=fallback_engine,
)

# Response caching (in-memory; swap for Redis in prod)
from llama_index.core.query_engine import CachingQueryEngine
cached_engine = CachingQueryEngine(retry_engine)

# Use cached_engine everywhere
response = cached_engine.query("Rate limit policy?")

Verify: Simulate a timeout (TimeoutException) and confirm retry kicks in. Hit the same query twice — second call should return instantly from cache (check logs for “Cache hit”).

Step 10: Deploy the query engine as a service

Wrap the engine in a stateless API. This pattern scales horizontally.

# app/main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from contextlib import asynccontextmanager

from llama_index.core import VectorStoreIndex, Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

class QueryRequest(BaseModel):
    question: str
    top_k: int = 4
    stream: bool = False

class QueryResponse(BaseModel):
    answer: str
    sources: list[dict]

index: VectorStoreIndex | None = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global index
    Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0.2)
    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
    
    chroma = chromadb.HttpClient(host="chroma", port=8000)
    vector_store = ChromaVectorStore(chroma_collection=chroma.get_or_create_collection("docs"))
    index = VectorStoreIndex.from_vector_store(vector_store)
    yield

app = FastAPI(lifespan=lifespan)

@app.post("/query", response_model=QueryResponse)
async def query(req: QueryRequest):
    if index is None:
        raise HTTPException(503, "Index not ready")
    
    engine = index.as_query_engine(
        similarity_top_k=req.top_k,
        response_mode="compact",
        streaming=req.stream,
    )
    
    if req.stream:
        # Return StreamingResponse instead
        from fastapi.responses import StreamingResponse
        async def gen():
            resp = await engine.aquery(req.question)
            async for token in resp.async_response_gen():
                yield token
        return StreamingResponse(gen(), media_type="text/plain")
    
    response = engine.query(req.question)
    return QueryResponse(
        answer=response.response,
        sources=[{"text": n.text[:200], "score": n.score} for n in response.source_nodes],
    )

Verify: Deploy to staging. Run curl -X POST /query -d '{"question":"rate limit"}'. Confirm JSON response with answer and sources. Load test with hey or locust — latency should be dominated by LLM + retrieval, not framework overhead.

What you’ve replaced

Common migration pitfalls

Chunk size mismatch: If you embedded with 512-token chunks but query with 1024-token context windows, you’ll waste tokens or truncate. Set Settings.chunk_size and Settings.chunk_overlap consistently at index time.

Embedding model drift: Re-embed if you switch from text-embedding-ada-002 to text-embedding-3-small. Mixed embeddings in one index break retrieval.

Sync vs async mixing: Don’t call .query() inside an async endpoint. Use .aquery() and async_response_gen. The reverse blocks the event loop.

Global Settings pollution: Settings.llm is process-global. In multi-tenant apps, pass llm= and embed_model= explicitly to each component instead.

Next steps

• Add a router query engine to route questions to specialized indexes (SQL, vector, keyword).
• Implement query transformation (HyDE, step-back prompting) for harder questions.
• Set up continuous evaluation in CI — fail deploys when faithfulness drops below threshold.
• If you run multiple model providers, swap OpenAI for Anthropic, Gemini, or a unified gateway like n4n.ai that handles fallback and metering across 240+ models behind one OpenAI-compatible endpoint.

The migration is complete when you can change the retrieval strategy, swap the LLM, or add a new evaluator without touching your application logic. That’s the architecture LlamaIndex gives you.