LlamaIndex — Testing & Production¶

import hashlib
import math
import re

from llama_index.core.base.embeddings.base import BaseEmbedding


class KeywordEmbedding(BaseEmbedding):
    """Hashed bag of words: same text, same vector; shared words, higher cosine similarity."""

    dim: int = 128

    def _vec(self, text: str) -> list[float]:
        vec = [0.0] * self.dim
        for token in re.findall(r"[a-z0-9]+", text.lower()):
            vec[int(hashlib.md5(token.encode()).hexdigest(), 16) % self.dim] += 1.0
        norm = math.sqrt(sum(x * x for x in vec)) or 1.0
        return [x / norm for x in vec]

    def _get_query_embedding(self, query: str) -> list[float]:
        return self._vec(query)

    async def _aget_query_embedding(self, query: str) -> list[float]:
        return self._vec(query)

    def _get_text_embedding(self, text: str) -> list[float]:
        return self._vec(text)

("ship-express", "Express shipping takes 1 to 2 business days and costs 15 USD.", "shipping", "public"),
("sec-incident", "Internal runbook: rotate API keys within one hour of a suspected leak.", "security", "internal"),

import pytest
from llama_index.core import VectorStoreIndex
from llama_index.core.llms import MockLLM
from llama_index.core.node_parser import SentenceSplitter

from corpus import make_docs
from kwembed import KeywordEmbedding


@pytest.fixture(autouse=True)
def no_real_provider(monkeypatch):
    """A test that forgets to inject a model fails on the missing key instead of calling OpenAI."""
    monkeypatch.delenv("OPENAI_API_KEY", raising=False)


@pytest.fixture
def splitter():
    return SentenceSplitter(chunk_size=256, chunk_overlap=16, id_func=lambda i, doc: f"{doc.doc_id}#{i}")


@pytest.fixture
def index(splitter):
    return VectorStoreIndex.from_documents(make_docs(), embed_model=KeywordEmbedding(), transformations=[splitter])


@pytest.fixture
def llm():
    return MockLLM()                     # complete() returns the prompt it received

import pytest
from llama_index.core import Document
from llama_index.core.vector_stores import FilterOperator, MetadataFilter, MetadataFilters


@pytest.mark.parametrize(
    ("query", "expected_id"),
    [
        ("express delivery takes 1 to 2 days", "ship-express#0"),
        ("digital goods refundable if download link never used", "refund-digital#0"),
        ("minimum password length", "sec-password#0"),
    ],
)
def test_top1_is_the_expected_chunk(index, query, expected_id):
    nodes = index.as_retriever(similarity_top_k=1).retrieve(query)
    assert [n.node_id for n in nodes] == [expected_id]
    assert nodes[0].score > 0.1


def test_acl_filter_never_leaks_internal_chunks(index):
    public = MetadataFilters(filters=[MetadataFilter(key="visibility", value="public", operator=FilterOperator.EQ)])
    nodes = index.as_retriever(similarity_top_k=5, filters=public).retrieve("rotate API keys after a leak")
    assert nodes and all(n.metadata["visibility"] == "public" for n in nodes)
    assert "sec-incident#0" not in {n.node_id for n in nodes}


def test_chunking_snapshot(splitter):
    text = " ".join(f"Sentence number {i} is here." for i in range(40))
    nodes = splitter.get_nodes_from_documents([Document(text=text, doc_id="doc")])
    assert [n.node_id for n in nodes] == ["doc#0", "doc#1"]
    assert nodes[0].text.startswith("Sentence number 0 is here.") and nodes[1].ref_doc_id == "doc"


def test_prompt_contains_retrieved_context(index, llm):
    response = index.as_query_engine(llm=llm, similarity_top_k=1).query("express delivery takes 1 to 2 days")
    assert "Express shipping takes 1 to 2 business days" in str(response)      # MockLLM echoes the prompt
    assert [n.node_id for n in response.source_nodes] == ["ship-express#0"]

from llama_index.core.agent.workflow import FunctionAgent, ToolCallResult
from llama_index.core.base.llms.types import ChatMessage, MessageRole, ToolCallBlock
from llama_index.core.llms import MockFunctionCallingLLM
from llama_index.core.tools import QueryEngineTool


def scripted_llm(*steps):
    """Each step is ("tool_name", {kwargs}) or a final answer string, consumed one per LLM call."""
    queue = list(steps)

    def generator(messages, **kwargs):
        current = queue.pop(0)
        if isinstance(current, str):
            return ChatMessage(role=MessageRole.ASSISTANT, content=current)
        name, args = current
        call = ToolCallBlock(tool_call_id="call-1", tool_name=name, tool_kwargs=args)
        return ChatMessage(role=MessageRole.ASSISTANT, blocks=[call])

    return MockFunctionCallingLLM(response_generator=generator)


async def test_agent_calls_the_knowledge_base_tool(index, llm):
    kb = QueryEngineTool.from_defaults(index.as_query_engine(llm=llm, similarity_top_k=1), name="kb",
                                       description="Search the support knowledge base.")
    llm_script = scripted_llm(("kb", {"input": "express delivery takes 1 to 2 days"}), "Express takes 1 to 2 days.")
    handler = FunctionAgent(tools=[kb], llm=llm_script).run("How fast is express shipping?")
    results = [e async for e in handler.stream_events() if isinstance(e, ToolCallResult)]
    assert str(await handler) == "Express takes 1 to 2 days."
    assert [r.tool_name for r in results] == ["kb"]
    assert "Express shipping takes 1 to 2 business days" in results[0].tool_output.content

from rag_flow import RAGWorkflow


async def test_workflow_paths_and_graph(index, llm):
    flow = RAGWorkflow(index, llm, timeout=10)
    response = await flow.run(query="express delivery takes 1 to 2 days")
    assert response.source_nodes[0].node_id == "ship-express#0"
    assert await flow.run(query="") == "Please ask a question."
    flow.validate()                                     # every event has a producer and a consumer

from llama_index.core.agent.workflow import FunctionAgent


async def test_tool_call_round_trip_over_http(openai_llm, fake_openai):
    def order_status(order_id: str) -> str:
        """Return the status of an order."""
        return "shipped"

    fake_openai.script[:] = [{"tool": "order_status", "args": {"order_id": "o1"}}, "Order o1 is shipped."]
    answer = await FunctionAgent(tools=[order_status], llm=openai_llm).run("Where is o1?")
    first, second = (r["body"] for r in fake_openai.requests)
    assert str(answer).strip() == "Order o1 is shipped."
    assert first["tools"][0]["function"]["name"] == "order_status"
    assert [m["content"] for m in second["messages"] if m["role"] == "tool"] == ["shipped"]

from llama_index.core.evaluation import RetrieverEvaluator

GOLDEN = [  # (question, expected node ids): reviewed by a human, frozen in git
    ("express delivery takes 1 to 2 days", ["ship-express#0"]),
    ("standard shipping free for orders above 50 USD", ["ship-standard#0"]),
    ("digital goods refundable if download link never used", ["refund-digital#0"]),
    ("passwords must contain at least 12 characters", ["sec-password#0"]),
]


def test_retrieval_quality_gate(index):
    evaluator = RetrieverEvaluator.from_metric_names(["hit_rate", "mrr"], retriever=index.as_retriever(similarity_top_k=3))
    results = [evaluator.evaluate(query=q, expected_ids=ids) for q, ids in GOLDEN]
    for metric, floor in {"hit_rate": 0.95, "mrr": 0.9}.items():       # baseline 1.0 minus tolerance
        assert sum(r.metric_vals_dict[metric] for r in results) / len(results) >= floor, metric

@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="real judge needs a key")
async def test_answers_are_faithful(real_query_engine, judge):
    runner = BatchEvalRunner({"faithfulness": FaithfulnessEvaluator(llm=judge)}, workers=4)
    results = await runner.aevaluate_queries(real_query_engine, queries=[q for q, _ in GOLDEN])
    assert sum(r.passing for r in results["faithfulness"]) / len(GOLDEN) >= 0.9

from llama_index.core.rate_limiter import TokenBucketRateLimiter
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

limiter = TokenBucketRateLimiter(requests_per_minute=300, tokens_per_minute=200_000)    # one shared budget
llm = OpenAI(model="gpt-4o-mini", rate_limiter=limiter, max_retries=3, timeout=30)
embed = OpenAIEmbedding(model="text-embedding-3-small", rate_limiter=limiter, embed_batch_size=100, num_workers=4)

def filters_for(role: str) -> MetadataFilters | None:
    """Built from the caller's identity, never from the request body."""
    if role == "staff":
        return None
    return MetadataFilters(filters=[MetadataFilter(key="visibility", value="public", operator=FilterOperator.EQ)])


@app.post("/ask")
async def ask(question: str, role: str = Depends(current_role)) -> dict:     # current_role: your auth dependency
    engine = index.as_query_engine(llm=llm, similarity_top_k=4, filters=filters_for(role))
    response = await engine.aquery(question)
    return {"answer": str(response), "sources": [{"node_id": n.node_id, "score": n.score} for n in response.source_nodes]}