BYO Vector Store: LlamaCloud Index V2 → Azure AI Search

%%capture
# Installed unconditionally rather than behind a `try: import ... except ImportError`
# guard: a different revision already present in the environment would otherwise
# shadow the pinned commit, and you'd be reading results this recipe never verified.
# `--upgrade` makes pip re-resolve the direct git reference instead of keeping
# whatever is already installed. The pinned commit declares requires-python ">=3.11,<3.14".
%pip install --quiet --upgrade "index-v2-data-sinks@git+https://github.com/run-llama/index-v2-data-sinks@798cac7e25f7485aa78b5f666957af41cae62c2f" "azure-search-documents==12.0.0" "azure-identity==1.25.3" "openai==1.109.1" "beautifulsoup4==4.13.3" "python-dotenv==1.2.2"

# --- Configuration ---------------------------------------------------------
# Index naming. The sink provisions ONE index per embedding model; the default
# model keeps the base name, others get a sanitized suffix.
#
# The prefix carries a short random suffix so this demo always creates its OWN
# index. A fixed, predictable name would let this notebook delete an index that
# an earlier run — or a colleague on the same service — is still using.
import uuid

RUN_ID = uuid.uuid4().hex[:8]
INDEX_PREFIX = f"forgebook-index-v2-sink-{RUN_ID}"

# Which model the sink sizes the index for. The Azure OpenAI deployment configured
# below must be the same base model, or query-time vectors and the index dimension will diverge.
EMBEDDING_MODELS = {
    "openai:text-embedding-3-large": 3072,
    "openai:text-embedding-3-small": 1536,
    "openai:text-embedding-ada-002": 1536,
}
EMBEDDING_MODEL = "openai:text-embedding-3-large"  # Change this with your deployment.
EMBEDDING_DIM = EMBEDDING_MODELS[EMBEDDING_MODEL]
AOAI_API_VERSION = "2024-10-21"   # Azure OpenAI data-plane api-version
EMBED_BATCH = 16                  # chunks embedded per Azure OpenAI call


def sink_index_name(prefix, embedding_model):
    """Match the official sink's one-index-per-model naming convention."""
    if embedding_model == "openai:text-embedding-3-small":
        return prefix
    suffix = embedding_model.replace(":", "-").replace("/", "-").replace(".", "-").replace("_", "-").lower()
    return f"{prefix}-{suffix}"


INDEX_NAME = sink_index_name(INDEX_PREFIX, EMBEDDING_MODEL)

# How the sink chunks each document: "fast" | "sentence" | "token".
CHUNK_STRATEGY = "sentence"

# Tenant + export-config labels written onto every chunk. In production these are
# your LlamaCloud project id and export-config id; here they're the scope fields
# queries can filter on, and the scope the sink uses for per-file replace and snapshots.
PROJECT_ID = "forgebook-demo"     # -> tenant_id on every document
CONFIG_ID = "forgebook-config"    # -> export_config_id; the snapshot/delete scope

# Retrieval Harness knobs.
TOP_K = 5        # final hits returned by Hybrid Retrieve
K_NEAREST = 50   # vector candidates fused before ranking; raise for recall
GREP_TOP = 50    # max chunks File Grep scans in one file
READ_TOP = 1000  # whole-file ceiling for File Read
LIST_TOP = 1000  # max docs List Files scans to build the map

# The demo corpus: three Paul Graham essays, each becomes one "file".
ESSAYS = {
    "great-work": "https://paulgraham.com/greatwork.html",
    "superlinear": "https://paulgraham.com/superlinear.html",
    "how-to-do-what-you-love": "https://paulgraham.com/love.html",
}
# ---------------------------------------------------------------------------
print(f"Config set. Index '{INDEX_NAME}' (unique to this run), {len(ESSAYS)} source files, top_k={TOP_K}.")

import os

from azure.core.credentials import AzureKeyCredential
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
from dotenv import load_dotenv
from openai import AzureOpenAI

load_dotenv(override=True)


def env(name, *, required=True, default=None):
    value = os.getenv(name, default)
    if required and not value:
        raise RuntimeError(f"Missing required env var: {name}. See Prerequisites.")
    return value or None


SEARCH_ENDPOINT = env("SEARCH_ENDPOINT")
AOAI_ENDPOINT = env("AOAI_ENDPOINT").split("/openai/", 1)[0].rstrip("/")
EMBED_DEPLOYMENT = env(
    "AOAI_EMBEDDING_DEPLOYMENT",
    required=False,
    default="text-embedding-3-large",
)  # Use the same base model as EMBEDDING_MODEL.

# Optional key fallback — leave both unset for the keyless default path.
SEARCH_API_KEY = env("SEARCH_API_KEY", required=False)
AOAI_API_KEY = env("AOAI_API_KEY", required=False)

# The sink's AzureSearchService reads its target from these env vars. With no
# API key set, it resolves DefaultAzureCredential — the same identity as ours.
os.environ["DEFAULT_INDEX_AZURE_SEARCH_ENDPOINT"] = SEARCH_ENDPOINT
os.environ["DEFAULT_INDEX_AZURE_SEARCH_INDEX_PREFIX"] = INDEX_PREFIX
if SEARCH_API_KEY:
    os.environ["DEFAULT_INDEX_AZURE_SEARCH_API_KEY"] = SEARCH_API_KEY

# Credential for the clients this notebook creates directly.
search_credential = AzureKeyCredential(SEARCH_API_KEY) if SEARCH_API_KEY else DefaultAzureCredential()

if AOAI_API_KEY:
    aoai = AzureOpenAI(
        azure_endpoint=AOAI_ENDPOINT,
        api_key=AOAI_API_KEY,
        api_version=AOAI_API_VERSION,
    )
else:
    aoai = AzureOpenAI(
        azure_endpoint=AOAI_ENDPOINT,
        azure_ad_token_provider=get_bearer_token_provider(
            DefaultAzureCredential(),
            "https://cognitiveservices.azure.com/.default",
        ),
        api_version=AOAI_API_VERSION,
    )


def _escape_odata(value):
    """OData string literals escape embedded single quotes by doubling them.
    Anything interpolated into a filter goes through this."""
    return value.replace("'", "''")


def embed(texts):
    """Embed with your Azure OpenAI deployment, batched."""
    vectors = []
    for i in range(0, len(texts), EMBED_BATCH):
        response = aoai.embeddings.create(
            model=EMBED_DEPLOYMENT,
            input=texts[i : i + EMBED_BATCH],
        )
        vectors.extend(item.embedding for item in response.data)
    return vectors


auth_mode = "API key (fallback)" if SEARCH_API_KEY else "Entra ID (keyless)"
print(f"Auth: {auth_mode} | embeddings: {EMBED_DEPLOYMENT} | sink will provision: {INDEX_NAME}")

import importlib.metadata as md
import json
import sys

# --- The environment is what we pinned -------------------------------------
PINNED_SHA = "798cac7e25f7485aa78b5f666957af41cae62c2f"
INSTALL_CMD = (
    '%pip install --quiet --upgrade '
    '"index-v2-data-sinks@git+https://github.com/run-llama/index-v2-data-sinks@' + PINNED_SHA + '" '
    '"azure-search-documents==12.0.0" "azure-identity==1.25.3"'
)

assert (3, 11) <= sys.version_info[:2] < (3, 14), (
    f"index-v2-data-sinks requires Python >=3.11,<3.14; this kernel is "
    f"{sys.version_info.major}.{sys.version_info.minor}. Switch kernels and re-run."
)

sink_dist = md.distribution("index-v2-data-sinks")
sink_sha = (json.loads(sink_dist.read_text("direct_url.json") or "{}")
            .get("vcs_info", {}).get("commit_id"))
assert sink_sha == PINNED_SHA, (
    f"index-v2-data-sinks is at {sink_sha}, expected {PINNED_SHA}. "
    f"Re-run the setup cell, or:\n    {INSTALL_CMD}"
)

for pkg, want in [("azure-search-documents", "12.0.0"), ("azure-identity", "1.25.3")]:
    got = md.version(pkg)
    assert got == want, f"{pkg} is {got}, expected {want}. Fix with:\n    {INSTALL_CMD}"

# --- The credential and the embedding deployment work ----------------------
from azure.search.documents.indexes import SearchIndexClient

probe = embed(["setup probe"])[0]
assert len(probe) == EMBEDDING_DIM, (
    f"Embedding returned {len(probe)} dims, expected {EMBEDDING_DIM} for {EMBEDDING_MODEL} — "
    f"is AOAI_EMBEDDING_DEPLOYMENT a deployment of that model?"
)

index_client = SearchIndexClient(SEARCH_ENDPOINT, search_credential)
existing = list(index_client.list_index_names())
assert INDEX_NAME not in existing, (
    f"'{INDEX_NAME}' already exists. The run id is random, so this is a collision — "
    f"re-run the Configuration cell to draw a new one."
)

print(f"Setup OK — sink @ {sink_sha[:7]}, Python {sys.version_info.major}.{sys.version_info.minor}, "
      f"embeddings return {len(probe)} dims.")
print(f"Azure AI Search reachable ({len(existing)} indexes); '{INDEX_NAME}' is free to create.")

import re
import urllib.request

from bs4 import BeautifulSoup
from index_v2_data_sinks.chunking import ChunkerService
from index_v2_data_sinks.schema import ExportedOutputV1, ParsedPage, ParseSyncOutputV1


def fetch_text(url):
    request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(request, timeout=30) as response:
        html = response.read().decode("utf-8", errors="ignore")
    return re.sub(r"\n{3,}", "\n\n", BeautifulSoup(html, "html.parser").get_text("\n")).strip()


chunker = ChunkerService()
files = {}  # parsed_directory_file_id -> list[ExportedOutputV1]
for file_id, url in ESSAYS.items():
    page_text = fetch_text(url)
    parse_output = ParseSyncOutputV1(pages=[ParsedPage(page_number=1, text=page_text)])
    # Single-page docs, so the chunker input is the page text itself.
    # ExportedOutputV1.build maps each chunk back to its page range internally.
    chunks = chunker.chunk(page_text, CHUNK_STRATEGY)
    files[file_id] = [
        ExportedOutputV1.build(parse_output, PROJECT_ID, CONFIG_ID, chunk, file_id, metadata={"file_name": file_id})
        for chunk in chunks
    ]
    print(f"{file_id:>26}: {len(files[file_id]):3d} chunks")

from index_v2_data_sinks.export.azure_search import AzureSearchExporter

exporter = AzureSearchExporter()

# Record what this run creates so the teardown cell deletes only this, and only
# if the export actually got that far.
CREATED_INDEXES = set()

total = 0
for file_id, outputs in files.items():
    embeddings = embed([o.content for o in outputs])
    assert len(embeddings) == len(outputs), (
        f"{file_id}: {len(embeddings)} embeddings for {len(outputs)} chunks — these must be 1:1."
    )
    await exporter.export(
        outputs,
        project_id=PROJECT_ID,
        embeddings=embeddings,
        embedding_model=EMBEDDING_MODEL,
    )
    CREATED_INDEXES.add(INDEX_NAME)
    total += len(outputs)
    print(f"exported {file_id}: {len(outputs)} chunks")

expected_total = sum(len(o) for o in files.values())
assert total == expected_total, f"Exported {total} chunks, expected {expected_total}."
assert INDEX_NAME in index_client.list_index_names(), (
    f"The exporter did not create '{INDEX_NAME}'."
)

print(f"\nExported {total} chunks via the official AzureSearchExporter into '{INDEX_NAME}'.")
print(f"Created by this run (teardown will delete only this): {sorted(CREATED_INDEXES)}")

import time

from azure.search.documents import SearchClient

search_client = SearchClient(SEARCH_ENDPOINT, INDEX_NAME, search_credential)


def wait_for_docs(client, target, timeout=30):
    """Azure AI Search indexing is near-real-time, not instant — poll until
    the document count reaches the target (or the timeout passes)."""
    for _ in range(timeout):
        if client.get_document_count() >= target:
            break
        time.sleep(1)
    return client.get_document_count()


count = wait_for_docs(search_client, total)
assert count == total, (
    f"Index holds {count} documents, expected {total}. Azure AI Search indexing is "
    f"near-real-time — if this is short, raise the wait_for_docs timeout and re-run."
)

schema = index_client.get_index(INDEX_NAME)
field_names = {f.name for f in schema.fields}
for required in ("id", "content", "embedding", "tenant_id", "export_config_id", "parsed_directory_file_id", "chunk_index"):
    assert required in field_names, f"The sink did not provision the '{required}' field."

vector_field = next(f for f in schema.fields if f.name == "embedding")
assert vector_field.vector_search_dimensions == EMBEDDING_DIM, (
    f"Index vector field is {vector_field.vector_search_dimensions}-dim, but {EMBEDDING_MODEL} "
    f"produces {EMBEDDING_DIM}-dim vectors."
)
print("Index fields:", [f.name for f in schema.fields])
print("Documents:", count)

doc = next(iter(search_client.search(
    search_text="*",
    select=["id", "parsed_directory_file_id", "tenant_id", "chunk_index", "content"],
)))
print("\nRaw document:")
print("  id:", doc["id"])
print("  tenant_id:", doc["tenant_id"], "| file:", doc["parsed_directory_file_id"], "| chunk_index:", doc["chunk_index"])
print("  content:", " ".join(doc["content"].split())[:160], "...")

same_query = "What does it take to do great work?"


def scoped(tenant):
    return list(search_client.search(
        search_text=same_query,
        filter=f"tenant_id eq '{_escape_odata(tenant)}'",
        select=["parsed_directory_file_id", "chunk_index", "tenant_id"],
        top=TOP_K,
    ))


mine = scoped(PROJECT_ID)
other = scoped("another-tenant")
unscoped = list(search_client.search(search_text=same_query, select=["tenant_id"], top=TOP_K))

assert mine, f"Expected hits for tenant '{PROJECT_ID}'."
assert all(r["tenant_id"] == PROJECT_ID for r in mine), "A scoped query returned an out-of-scope document."
assert not other, f"Expected 0 hits for an unused scope, got {len(other)}."
assert unscoped, "Expected the unfiltered query to return documents."

print(f"scope '{PROJECT_ID}': {len(mine)} hits  |  scope 'another-tenant': {len(other)} hits")
print(f"no filter at all:     {len(unscoped)} hits  <- the same caller, unscoped, sees everything")

from azure.search.documents.models import VectorizedQuery


class SinkRetrievalHarness:
    """Retrieval Harness primitives over an index-v2-data-sinks Azure AI Search index."""

    def __init__(self, client, embed_fn):
        self.client = client
        self.embed_fn = embed_fn
        self.select = ["parsed_directory_file_id", "chunk_index", "content"]

    def hybrid_retrieve(self, query, top_k=TOP_K, k_nearest=K_NEAREST):
        """High-recall first pass: vector + BM25 fusion over the sink index."""
        vector = self.embed_fn([query])[0]
        results = self.client.search(
            search_text=query,
            vector_queries=[VectorizedQuery(vector=vector, k_nearest_neighbors=k_nearest, fields="embedding")],
            select=self.select,
            top=top_k,
        )
        return [self._row(r) for r in results]

    def list_files(self, top=LIST_TOP):
        """Map the corpus: every parsed_directory_file_id and its chunk count."""
        results = self.client.search(search_text="*", select=["parsed_directory_file_id"], top=top)
        counts = {}
        for r in results:
            counts[r["parsed_directory_file_id"]] = counts.get(r["parsed_directory_file_id"], 0) + 1
        return [{"file": name, "chunks": n} for name, n in sorted(counts.items())]

    def file_grep(self, file_id, term, top=GREP_TOP):
        """Server-side term scan of one file. Matching goes through the query
        analyzer — tokenized and case-insensitive, not byte-exact."""
        odata = (
            f"parsed_directory_file_id eq '{_escape_odata(file_id)}' and "
            f"search.ismatch('{_escape_odata(term)}', 'content', 'full', 'any')"
        )
        results = self.client.search(search_text="*", filter=odata, select=self.select, top=top)
        return sorted((self._row(r) for r in results), key=lambda d: d["chunk_index"])

    def file_read(self, file_id, top=READ_TOP):
        """Pull a whole file back in chunk order."""
        results = self.client.search(
            search_text="*",
            filter=f"parsed_directory_file_id eq '{_escape_odata(file_id)}'",
            select=self.select,
            top=top,
        )
        return sorted((self._row(r) for r in results), key=lambda d: d["chunk_index"])

    def _row(self, r):
        return {
            "file": r.get("parsed_directory_file_id"),
            "chunk_index": r.get("chunk_index"),
            "score": r.get("@search.score"),
            "text": r.get("content"),
        }


harness = SinkRetrievalHarness(search_client, embed)


def show(rows, n=200):
    for r in rows:
        head = f"[{r['file']} #{r['chunk_index']:>3}]"
        if r.get("score") is not None:
            head += f"  score={r['score']:.4f}"
        print(head)
        print("  ", " ".join((r.get("text") or "").split())[:n], "...\n")


print("Harness ready:", ["hybrid_retrieve", "list_files", "file_grep", "file_read"])

answers = harness.hybrid_retrieve("What does it take to do great work?")

assert len(answers) == TOP_K, f"Hybrid Retrieve returned {len(answers)} hits, expected top_k={TOP_K}."
assert all({"file", "chunk_index", "score", "text"} <= r.keys() for r in answers), "Unexpected row shape."
assert all(r["score"] is not None and r["text"] for r in answers), "Every hit should carry a score and text."
assert answers == sorted(answers, key=lambda r: -r["score"]), "Hits should arrive ranked."

show(answers)

listing = harness.list_files()

assert {f["file"] for f in listing} == set(ESSAYS), (
    f"List Files saw {sorted(f['file'] for f in listing)}, expected {sorted(ESSAYS)}."
)
assert sum(f["chunks"] for f in listing) == total, (
    f"Chunk counts sum to {sum(f['chunks'] for f in listing)}, expected {total}."
)

for f in listing:
    print(f"{f['file']:>26}  ({f['chunks']} chunks)")

TERM, TARGET_FILE = "exponential", "superlinear"
hits = harness.file_grep(TARGET_FILE, TERM)

assert hits, f"Expected at least one '{TERM}' hit in {TARGET_FILE}."
assert all(h["file"] == TARGET_FILE for h in hits), "File Grep leaked chunks from another file."
assert all(TERM in (h["text"] or "").lower() for h in hits), (
    f"Every hit should contain '{TERM}' — the analyzer is case-insensitive, so compare lowercased."
)
assert not harness.file_grep(TARGET_FILE, "zzzznotarealterm"), "A nonsense term should return nothing."

print(f"'{TERM}' -> {len(hits)} chunk(s) in {TARGET_FILE}; a nonsense term -> 0\n")
show(hits[:3])

TARGET = "how-to-do-what-you-love"
chapter = harness.file_read(TARGET)

expected_chunks = next(f["chunks"] for f in listing if f["file"] == TARGET)
assert len(chapter) == expected_chunks, (
    f"File Read returned {len(chapter)} chunks, but List Files counted {expected_chunks}."
)
assert all(r["file"] == TARGET for r in chapter), "File Read leaked chunks from another file."
indexes = [r["chunk_index"] for r in chapter]
assert indexes == sorted(indexes), "File Read must return chunks in order."

print(f"Reassembled {len(chapter)} chunks in order (indexes {indexes[0]}..{indexes[-1]}):\n")
for r in chapter[:3]:
    print(f"#{r['chunk_index']:>3}  {' '.join((r['text'] or '').split())[:200]}...")

before = search_client.get_document_count()

again = files["superlinear"]
await exporter.export(
    again, project_id=PROJECT_ID, embeddings=embed([o.content for o in again]), embedding_model=EMBEDDING_MODEL
)

# merge_or_upload with deterministic keys means the count must NOT grow. Poll for
# a few seconds so a would-be duplicate has time to surface before we compare —
# a count that is still equal only proves idempotency once indexing has settled.
deadline = time.time() + 8
while time.time() < deadline and search_client.get_document_count() <= before:
    time.sleep(1)

after = search_client.get_document_count()
assert after == before, (
    f"Re-export changed the document count {before} -> {after}. Deterministic keys plus "
    f"merge_or_upload should make a re-export a no-op; a growing count means duplicate keys."
)
print(f"doc count before re-export: {before}  |  after: {after}  ->  no duplicates")

snapshots = await exporter.list_snapshots(export_config_id=CONFIG_ID, embedding_model=EMBEDDING_MODEL)
snapshot_files = sorted(s.parsed_directory_file_id for s in snapshots)
assert snapshot_files == sorted(ESSAYS), (
    f"Expected one snapshot per file {sorted(ESSAYS)}, got {snapshot_files}."
)
print("snapshots (one per file):", snapshot_files)

DELETE_INDEX = True

if not DELETE_INDEX:
    print(f"Kept {sorted(CREATED_INDEXES)} — inspect in Azure portal > Search service > Indexes.")
else:
    for name in sorted(CREATED_INDEXES):
        index_client.delete_index(name)
        print(f"Deleted index '{name}'.")
    remaining = set(index_client.list_index_names())
    assert not (CREATED_INDEXES & remaining), "Teardown left an index behind."
    print("Teardown complete — only indexes created by this run were deleted.")