DEV Community: Divy Yadav

Multi-Aspect E-Commerce Semantic Engine Using Qdrant Multivectors

Divy Yadav — Wed, 01 Jul 2026 05:33:24 +0000

How I built a multi-vector semantic search engine that splits user intent before touching the database, using ColBERT, SigLIP, and BGE within a single Qdrant point.

Six months ago I was building a semantic search engine for a small e-commerce catalog. Everything looked fine. I was using a good embedding model, cosine similarity, and the search results seemed reasonable.

Then I searched for:

"Waterproof black hiking boots with good arch support."

The first result was perfect.

The second was a waterproof jacket.

The third was a pair of black Chelsea boots.

Technically, the search wasn't wrong. Every result was semantically similar to the query.

But if I were a customer, I wouldn't care. I wanted hiking boots with arch support, not products that happened to share a few similar words.

I tried a better embedding model. I changed the chunking strategy. I added more metadata. The results improved a little, but the real problem never went away.

The problem wasn't the model.

It was the assumption that one embedding could represent an entire product.

A product has technical specifications, images, and customer reviews. Those are different kinds of information, yet I was compressing all of them into a single vector.

So I rebuilt the system.

Instead of one embedding per product, I stored separate vector spaces for specifications, images, and review findings. I split each query into different types of intent before embedding it and added a lightweight personalization layer on top.

This article walks through how I built it, why I chose this architecture, and when this extra complexity is actually worth it.

The mental model: before you read a single line of code

Here's the full pipeline in one view. Read this once, keep it in mind, and the rest of the article will click much faster.

How to read this:

Think of the system as three specialists evaluating the same product.

When a user searches for "Waterproof black hiking boots with good arch support", the query is split into three questions:

Is it waterproof? → Specs expert (ColBERT)
Does it look like black hiking boots? → Visual expert (SigLIP)
Do customers praise the arch support? → Review expert (BGE)

Each expert searches a different vector field stored inside the same Qdrant point. Text and review retrieval nominate candidates first, the visual channel performs final scoring, and personalization reranks the results.

Instead of forcing one embedding to represent everything, the system evaluates independent pieces of evidence and combines them at ranking time.

The core problem: one vector can't represent everything

When someone searches "waterproof black hiking boots with good arch support," they're not asking one question. They're asking three at the same time:

The first is a spec question: "Is this waterproof?" That's answered by product descriptions and technical data sheets.

The second is a visual question: "Does it look like black hiking boots?" That's answered by images. Not text descriptions of images. Actual images.

The third is a social question: "Do people say the arch support is good?" That's answered by customer reviews.

No single embedding model can capture all three signals with equal fidelity. When you pool them together, you get a vector that's roughly in the neighborhood of all three but isn't precise about any of them.

It's why you get Chelsea boots in your hiking boot results. Both are black footwear. The visual signal is doing fine. The arch-support signal is completely lost.

The fix is to stop treating a product as one thing and start treating it as a collection of distinct signals, each searchable independently.

Why I chose Qdrant for this

I evaluated Pinecone, Weaviate, and Qdrant before committing. The deciding factor was named multivector fields.

I needed to store multiple different types of vectors per document, each with its own dimensionality and comparison function. Pinecone doesn't support this natively. Weaviate's multivector story was still evolving when I started.

Qdrant's API is clean: you define named vector configs at collection creation time, each with different dimensions, distance metrics, and HNSW settings.

The other thing that sold me was update_vectors. In a live catalog:

Products get new images. The visual matrix needs to grow without a collection rebuild.
Reviews arrive daily. Each new review adds rows to the review matrix.

Qdrant handles both by letting you retrieve an existing vector field, concatenate new vectors, and push the updated matrix back. The collection stays up. Running queries keep working.

Versions matter here. I'm on Qdrant v1.15.3 with qdrant-client>=1.15.0. Multivectors have been available since v1.10, but the query_points prefetch API that makes multi-stage search work cleanly became stable in v1.14.

What the system stores: one point, three matrices

Every product becomes a single Qdrant point. That point has three named vector fields:

# src/commerce_engine/qdrant_store.py

VISUAL_VECTOR = "visual_vectors"   # 768-d SigLIP image embedding
TEXT_VECTOR   = "text_vectors"     # 96-d ColBERT token matrix
REVIEW_VECTOR = "review_vectors"   # 384-d BGE per-finding embeddings

All three use MAX_SIM as the comparator. Here is the full collection setup:

def vector_params(size: int, *, hnsw_m: int | None = None) -> models.VectorParams:
    hnsw_config = None
    if hnsw_m is not None:
        hnsw_config = models.HnswConfigDiff(m=hnsw_m)
    return models.VectorParams(
        size=size,
        distance=models.Distance.COSINE,
        multivector_config=models.MultiVectorConfig(
            comparator=models.MultiVectorComparator.MAX_SIM
        ),
        hnsw_config=hnsw_config,
    )

def recreate_collection(client, collection, *, profile="baseline", disable_text_hnsw=True):
    client.create_collection(
        collection_name=collection,
        vectors_config={
            VISUAL_VECTOR: vector_params(VISION_DIM),                                       # 768
            TEXT_VECTOR:   vector_params(TEXT_DIM, hnsw_m=0 if disable_text_hnsw else None),# 96
            REVIEW_VECTOR: vector_params(REVIEW_DIM),                                       # 384
        },
        quantization_config=quantization_config(profile),
    )
    create_payload_indexes(client, collection)

The hnsw_m=0 on TEXT_VECTOR is intentional. Setting m=0 disables HNSW graph indexing entirely for that field.

Normally that would be a performance disaster. Here it isn't, for two reasons:

Text vectors are not used for first-stage retrieval. They're used as a ColBERT reranker over the small candidate set that the prefetch stages already returned. At 20–50 candidates, brute-force matrix comparison beats HNSW because the graph traversal overhead costs more than it saves at that N.
It saves index memory. Token-level matrices are larger than pooled vectors. Skipping the HNSW graph for text_vectors reduces memory footprint, which compounds when you're storing 96-d matrices across thousands of tokens per product.

The payload also gets indexed separately so Qdrant can filter before scoring:

def create_payload_indexes(client, collection):
    # Keyword fields
    for field in ["brand", "category", "region", "color", "size", "product_id"]:
        client.create_payload_index(
            collection_name=collection,
            field_name=field,
            field_schema=models.PayloadSchemaType.KEYWORD,
        )
    # Bool field
    client.create_payload_index(
        collection_name=collection,
        field_name="availability",
        field_schema=models.PayloadSchemaType.BOOL,
    )
    # Float fields
    for field in ["price", "eco_score"]:
        client.create_payload_index(
            collection_name=collection,
            field_name=field,
            field_schema=models.PayloadSchemaType.FLOAT,
        )

This matters more than it looks. When indexes are present, Qdrant runs payload filtering before vector scoring. Filtering availability=True or price <= 150.0 eliminates candidates before any embedding comparison runs. Without indexes you get post-filtering, which is slower and produces inconsistent result counts. Index your filterable fields before you query.

The three embedding pipelines

Each vector field in the point comes from a different embedding pipeline. The same product goes through all three before ingestion.

[Photo from AI]

Visual pipeline: SigLIP

Product images go through google/siglip-base-patch16-224.

SigLIP is a vision-language model with aligned image and text embeddings. Aligned means you can encode a text query like "black hiking boots" and compare it directly against an image embedding. No bridge model. No separate alignment step.

# src/commerce_engine/embeddings.py

def image_patches(self, image_path: Path) -> list[list[float]]:
    """Returns shape [1, 768]: normalized SigLIP pooled image embedding."""

    self._load_siglip()
    import torch
    image = Image.open(image_path).convert("RGB")
    siglip_inputs = self._siglip_processor(images=image, return_tensors="pt")
    siglip_inputs = {k: v.to(self.device) for k, v in siglip_inputs.items()}

    with torch.no_grad():
        vision_outputs = self._siglip_model.vision_model(**siglip_inputs)
        pooled = vision_outputs.pooler_output        # shape: [1, 768]
        normalized = torch.nn.functional.normalize(pooled, dim=-1)

    return normalized.cpu().numpy().astype(float).tolist()

And the corresponding text-side query encoding:

import torch

def visual_query(self, query: str) -> list[list[float]]:
    """Returns shape [1, 768]: normalized SigLIP text tower output."""
    self._load_siglip()

    text_inputs = self._siglip_processor(
        text=[query], return_tensors="pt", padding=True, truncation=True
    )

    text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}

    with torch.no_grad():
        text_outputs = self._siglip_model.text_model(**text_inputs)
        pooled = text_outputs.pooler_output          # shape: [1, 768]
        normalized = torch.nn.functional.normalize(pooled, dim=-1)

    return normalized.cpu().numpy().astype(float).tolist()

Both outputs are L2-normalized before storage and querying. This makes cosine similarity equivalent to dot product, which is faster for Qdrant to compute.

The stored shape for each product is [1, 768], a matrix with one row. It's technically a 2D multivector with one element, but storing it as a matrix rather than a flat vector keeps the interface consistent and makes appending additional image patches straightforward later.

Text pipeline: ColBERT

Product titles and spec sheets go through FastEmbed's answerdotai/answerai-colbert-small-v1. This model produces token-level embeddings, 96 dimensions per token, without pooling.

def text_late(self, texts: list[str]) -> list[list[list[float]]]:
    return [
        embedding.astype(float).tolist()
        for embedding in self.late_model.embed(texts)
    ]

The return shape is [num_texts, num_tokens, 96]. For one product title, you get a matrix of shape [num_tokens, 96]. A spec-heavy product like "TrailForge StormShield Black Hiking Boots. support: nylon shank and molded arch support. upper: black ripstop textile. waterproof_rating: IPX6 waterproof membrane for heavy rain" tokenizes to around 25-30 tokens, so the stored matrix is [~28, 96].

The text document is assembled in the Product model:

# src/commerce_engine/models.py

def text_document(self) -> str:
    specs = " ".join(f"{key}: {value}" for key, value in sorted(self.specs.items()))
    return f"{self.title}. {specs}"

Title first, then all specs as key-value pairs. The spec keys (support, waterproof_rating, upper) are included as text so ColBERT can match a query token like "waterproof" to the spec key, not just the spec value.

Review pipeline: BGE with semantic finding extraction

Raw reviews don't go straight to the embedder. They go through a finding extraction step first. This is where the architecture differs most from a naive approach.

A review like "Waterproof in heavy rain and mud. Good arch support on long hikes. Excellent grip on wet rock." holds three separate factual claims. Embedding the full review as one vector averages all three. A product with ten reviews embedded whole gives you ten averaged blobs. You lose the specific claim that matters to a specific query.

Instead, extract_semantic_findings() pulls out specific claims using regex patterns:

# src/commerce_engine/reviews.py

FINDING_PATTERNS = [
    r"waterproof(?: in [a-z ]+)?",
    r"kept feet dry",
    r"good arch support",
    r"arch support is [a-z]+",
    r"supportive footbed",
    r"excellent grip",
    r"not enough grip",
    r"highly durable",
    r"runs a little small",
    r"runs small",
    r"strong ankle support",
]

def extract_semantic_findings(reviews: list[str]) -> list[str]:
    findings: list[str] = []
    for review in reviews:
        normalized = review.lower().strip()
        matched = False
        for pattern in FINDING_PATTERNS:
            match = re.search(pattern, normalized)
            if match:
                findings.append(match.group(0))
                matched = True
        if not matched:
            findings.append(normalized.rstrip("."))
    return list(dict.fromkeys(findings))  # deduplication

The findings for the TrailForge boot come out as: ["waterproof in heavy rain and mud", "good arch support", "excellent grip"]. Each finding then gets embedded separately with BAAI/bge-small-en-v1.5 into a 384-dimensional vector:

def review_findings(self, findings: list[str]) -> list[list[float]]:
    if not findings:
        return [_unit_vector("empty-review", REVIEW_DIM)]
    return [
        embedding.astype(float).tolist()
        for embedding in self.review_model.embed(findings)
    ]

The stored review matrix is [num_findings, 384]. Three findings become a [3, 384] matrix. A product with ten reviews and 15 extracted findings becomes [15, 384].

When a query includes "good arch support," it gets embedded as a single 384-d vector and compared via MAX_SIM against this matrix. The most similar finding wins. That matching finding is "good arch support" from a review, not an approximation averaged across all review content.

The full ingestion step assembles all three matrices into a single point:

# src/commerce_engine/ingest.py

def product_point(product: Product, embedder: Embedder) -> models.PointStruct:
    text_matrix   = embedder.text_late([product.text_document()])[0]
    review_matrix = embedder.review_findings(
        extract_semantic_findings(product.reviews)
    )
    visual_matrix = embedder.image_patches(Path(product.image_path))

    return models.PointStruct(
        id=point_id(product.id),
        payload=product.payload(),
        vector={
            VISUAL_VECTOR: visual_matrix,   # shape [1, 768]
            TEXT_VECTOR:   text_matrix,     # shape [~28, 96]
            REVIEW_VECTOR: review_matrix,   # shape [num_findings, 384]
        },
    )

One product. One point. Three matrices, each solving a different retrieval problem.

Query decomposition: routing intent before embedding

When a user submits a query, the first step is decomposition. The query does not go directly to all three vector fields.

# src/commerce_engine/query.py

TEXT_KEYWORDS   = {"waterproof", "water-resistant", "rain", "membrane", "insulated", "leather"}
VISUAL_KEYWORDS = {"black", "brown", "charcoal", "red", "blue", "hiking", "boots", "sneakers"}
REVIEW_KEYWORDS = {"support", "arch", "grip", "durable", "comfortable", "runs", "excellent"}

def decompose_query(query: str) -> QueryPlan:
    tokens = tokenize(query)
    text_terms   = [t for t in tokens if t in TEXT_KEYWORDS]
    visual_terms = [t for t in tokens if t in VISUAL_KEYWORDS]
    review_terms = [t for t in tokens if t in REVIEW_KEYWORDS]
    phrase = " ".join(tokens)
    # phrase-level overrides for known compound terms
    if "good arch support" in phrase:
        review_terms.extend(["good", "arch", "support"])
    if "black hiking boots" in phrase:
        visual_terms.extend(["black", "hiking", "boots"])
    if "waterproof" in phrase:
        text_terms.append("waterproof")
    return QueryPlan(
        original_query=query,
        text_terms=list(dict.fromkeys(text_terms)),
        visual_terms=list(dict.fromkeys(visual_terms)),
        review_terms=list(dict.fromkeys(review_terms)),
    )

For "Waterproof black hiking boots with good arch support", the decomposition produces:

text_terms = ["waterproof"] → routed to ColBERT spec matching
visual_terms = ["black", "hiking", "boots"] → routed to SigLIP visual matching
review_terms = ["support", "arch", "good"] → routed to BGE review matching

The QueryPlan model exposes three properties that join these terms back into sub-query strings:

# src/commerce_engine/models.py

@property
def text_query(self) -> str:
    return " ".join(self.text_terms or [self.original_query])
@property
def visual_query(self) -> str:
    return " ".join(self.visual_terms or [self.original_query])
@property
def review_query(self) -> str:
    return " ".join(self.review_terms or [self.original_query])

If no terms match a category, the original query falls back as the sub-query for that field. This prevents empty queries.

I'll be direct about the limitation: this decomposition uses keyword matching, which is coarse. A better production implementation would call an LLM to classify query intent. The architecture is the same either way. What matters is that you route before you embed, not that you route using a fancy classifier. Get the structure right first.

The search: prefetch then final scoring

Qdrant's query_points supports a prefetch parameter. It lets you retrieve candidate sets from multiple vector fields in parallel, then apply a final scoring stage against only those candidates. This is the mechanism that ties the three vector fields together.

# src/commerce_engine/search.py

def search_products(client, collection, request, embedder):
    plan = decompose_query(request.query)
    query_filter = build_filter(request.filters)
    # Embed all three sub-queries
    text_query   = embedder.text_late([plan.text_query])[0]
    review_query = embedder.review_findings([plan.review_query])
    visual_query = embedder.visual_query(plan.visual_query)
    candidate_limit = max(request.limit * 5, 20)
    # Text and review run as prefetch candidates
    prefetch = [
        models.Prefetch(
            query=text_query,
            using=TEXT_VECTOR,
            limit=candidate_limit,
            filter=query_filter,
        ),
        models.Prefetch(
            query=review_query,
            using=REVIEW_VECTOR,
            limit=candidate_limit,
            filter=query_filter,
        ),
    ]

    # Visual query scores and ranks the prefetch candidates
    response = client.query_points(
        collection_name=collection,
        prefetch=prefetch,
        query=visual_query,
        using=VISUAL_VECTOR,
        query_filter=query_filter,
        limit=max(request.limit * 3, 10),
        with_payload=True,
        with_vectors=True,
    )

Text and review prefetch stages each return up to candidate_limit results. Qdrant unions them. The final visual query then scores only those candidates using the VISUAL_VECTOR field. Payload filters apply at every stage.

After Qdrant responds, the code computes per-aspect MAX_SIM scores to get a named breakdown:

from commerce_engine.scoring import maxsim_score

for point in response.points:
        vectors  = point.vector or {}
        doc_visual = vectors.get(VISUAL_VECTOR, [])
        doc_text   = vectors.get(TEXT_VECTOR, [])
        doc_review = vectors.get(REVIEW_VECTOR, [])
        v_score = maxsim_score(visual_query, doc_visual) if doc_visual else 0.0
        t_score = maxsim_score(text_query,   doc_text)   if doc_text   else 0.0
        r_score = maxsim_score(review_query, doc_review) if doc_review else 0.0
        combined_score = v_score + t_score + r_score

The MAX_SIM function itself is three lines of NumPy:

# src/commerce_engine/scoring.py

def maxsim_score(query_matrix, doc_matrix) -> float:
    query = np.asarray(query_matrix, dtype=np.float32)
    doc   = np.asarray(doc_matrix,   dtype=np.float32)
    similarities = query @ doc.T                # [query_tokens, doc_tokens]
    return float(similarities.max(axis=1).sum())# max per query token, then sum

For each query token: find the highest similarity against all document tokens. Sum those per-token maximums. That sum is the MAX_SIM score. This is the full ColBERT late-interaction formula.

Computing it manually after retrieval gives you per-aspect scores (visual=0.81, text=0.44, review=0.69) in the explanation field. That breakdown is the best debugging tool in the system. When a result feels wrong, the aspect scores tell you immediately whether the visual, spec, or review signal failed to match.

Personalization: reranking the same candidates differently

After Qdrant returns results, a personalization layer reranks them based on the requesting user's profile. The Qdrant retrieval step is identical for every user. Only the reranking changes.

# src/commerce_engine/scoring.py

def personalization_boost(payload: dict, profile: UserProfile) -> tuple[float, list[str]]:
    boost = 0.0
    reasons = []
    if payload.get("brand") in profile.preferred_brands:
        boost += 0.25
        reasons.append(f"preferred brand: {payload['brand']}")
    price = float(payload.get("price", 0.0))
    low, high = profile.price_range
    if low <= price <= high:
        boost += 0.30
        reasons.append(f"price in user range: {low:g}-{high:g}")
    if profile.eco_preference and payload.get("is_sustainable"):
        boost += 0.30
        reasons.append("eco preference matched")
    if payload.get("category") in profile.favorite_categories:
        boost += 0.20
        reasons.append(f"favorite category: {payload['category']}")
    return boost, reasons

def rerank(results: list[SearchResult], profile: UserProfile) -> list[SearchResult]:
    reranked = []
    for result in results:
        boost, reasons = personalization_boost(result.payload, profile)
        final_score = result.qdrant_score * (1.0 + boost)
        reranked.append(
            result.model_copy(update={
                "final_score": final_score,
                "personalization_boost": boost,
                "explanation": [*result.explanation, *reasons],
            })
        )
    return sorted(reranked, key=lambda r: r.final_score, reverse=True)

The boost is multiplicative, not additive. A product with a semantic score of 0.90 and a total boost of 0.55 finishes at 0.90 × 1.55 = 1.395. A product with a score of 0.95 and no matching preferences stays at 0.95. Personalization can flip the ranking.

Every matched boost reason gets appended to the result's explanation string, so a top result might show: "preferred brand: TrailForge | price in user range: 80-165 | review evidence: arch support". This matters for two reasons: users trust recommendations they can understand, and you can debug a bad result in seconds by looking at why its score was what it was.

Updating products without rebuilding anything

Production catalogs are not static. A product gets a new image. Reviews come in daily. The system needs to handle both without downtime.

Both update operations follow the same pattern: retrieve the existing vector field, concatenate new vectors, push the updated matrix back with update_vectors.

# src/commerce_engine/updates.py

def append_review(client, collection, product_id, review, embedder):
    point    = _get_point(client, collection, product_id)
    findings = extract_semantic_findings([review])
    new_vecs = embedder.review_findings(findings)
    existing        = (point.vector or {}).get(REVIEW_VECTOR, [])
    updated_matrix  = [*existing, *new_vecs]
    update_named_vectors(client, collection, product_id, {REVIEW_VECTOR: updated_matrix})
    client.set_payload(
        collection_name=collection,
        payload={"reviews": [*point.payload.get("reviews", []), review]},
        points=[point_id(product_id)],
        wait=True,
    )
    return {"product_id": product_id, "findings": findings}

def append_image(client, collection, product_id, image_path, embedder):
    point   = _get_point(client, collection, product_id)
    new_vecs = embedder.image_patches(image_path)
    existing       = (point.vector or {}).get(VISUAL_VECTOR, [])
    updated_matrix = [*existing, *new_vecs]
    update_named_vectors(client, collection, product_id, {VISUAL_VECTOR: updated_matrix})
    return {
        "product_id": product_id,
        "added_patches": len(new_vecs),
        "total_patches": len(updated_matrix),
    }

After append_review, the product's review matrix grows by one row per extracted finding. After append_image, the visual matrix gains one row. No collection rebuild. Running queries keep working.

This matters more than it sounds. If adding a single review required re-ingesting the whole product, you'd batch updates and your review data would always lag. With update_vectors, a webhook can process incoming reviews in real time, and the next search query immediately benefits from the new signal.

What the benchmarks actually showed

The benchmark runs 104 queries against the fixture product set (4 hiking boot variants), generated from product attribute combinations: color + category, brand + category, spec-based, review-fragment, and compound queries. Latencies measured with time.perf_counter against a live Qdrant instance.

Results for the baseline profile (no quantization, HNSW disabled for text):

95.2% Recall@3 means 99 of 104 queries returned the target product in the top 3. The 5 failures are edge cases where the query decomposer routes everything to one aspect but the product's signal lives in a different field.

Honest caveats worth stating:

This is 4 products. The 2.2ms will not hold at 100k products.
Benchmark queries were generated from the same fixture data they retrieve. Some overfitting is baked in.
Binary quantization knocked Recall@3 down noticeably. Compressing 32-bit floats to 1 bit is extreme for ColBERT's token-level matching. Start with scalar INT8 and measure recall before deploying.

When not to build this

Three embedding pipelines, a query decomposer, and prefetch orchestration is not always the right answer. Skip it if any of these apply:

Short, uniform queries like "red shoes" or "wool sweater." Single dense vector handles this fine.
No clean product images. SigLIP needs consistent product photography. Blurry or inconsistent shots produce weak visual signal. Drop that field.
No customer reviews. The review field adds complexity for zero gain. Skip it.
Catalog under 10k products with low query load. The simpler architecture will likely be fast enough and much cheaper to maintain.
Broken baseline retrieval. Bad base embeddings, wrong chunking, or misconfigured filters won't get fixed by adding more vectors. Multivectors amplify signal. If there's no signal to amplify, they amplify noise.

When this is worth building

The complexity pays off when all of these are true:

Queries contain mixed intent. Spec questions, visual questions, and review-based questions in the same search string. Common in outdoor gear, fashion, electronics, and beauty.
Clean product images at scale. 100k+ SKUs with consistent photography. The SigLIP channel earns its storage cost here.
Reviews arrive continuously. If reviews update daily, the update_vectors pattern becomes genuinely valuable versus batch re-indexing.
You have user profile data. The personalization layer needs signal to work with. Without it, the reranking step is adding noise.

Decision matrix

![Photo from AI

](https://dev-to-uploads.s3.us-east-2.amazonaws.com/uploads/articles/52ipdt5rv5se2mch6f2h.png)

Implementation path if you're starting from scratch

Do not skip step 1. The single-vector baseline is not just scaffolding. It is the control you need to measure whether the additional complexity actually improves results for your specific catalog.

I have seen teams add three embedding pipelines on day one because the architecture sounds impressive, then spend three weeks trying to figure out why recall got worse. It was worse because their base embeddings were bad. Fix the foundation before you add floors.

The thing that actually changed after building this

I went back and ran the original "Waterproof black hiking boots with good arch support" query through the finished system. Rank 1 was the TrailForge StormShield. Rank 2 was the EcoTrek TerraDry.

Both were actual hiking boots with documented arch support in customer reviews.

The Chelsea boots were gone.

The insight isn't that multivectors are magic. The insight is that the original setup was asking one number to represent three different kinds of evidence, and one number is not enough. The fix was to stop compressing and start separating.

Most search quality problems in production are not model problems. They are signal representation problems. The right model pointed at mixed-up data will keep returning wrong answers. Separating the signals is the work. The models are just the tools.

"Your search returns technically relevant results that nobody clicks? Ask yourself how many different types of user intent you are compressing into one similarity score."

The full source code, tests, benchmark runner, FastAPI endpoints, Streamlit UI, and real product dataset are all at:

Github

Everything described in this article maps to a real file in src/commerce_engine. No placeholder functions, no pseudo-code, no hand-waved implementation details.

Project at a glance

A quick summary of everything this system does, for anyone who wants the short version before going through the code.

What it is: A multi-aspect semantic search engine for e-commerce that splits every query into visual, spec, and review signals before retrieval.

Stack:

Vector database: Qdrant v1.15.3 with qdrant-client>=1.15.0
Visual embeddings: google/siglip-base-patch16-224 (768-d, via Hugging Face Transformers)
Text embeddings: answerdotai/answerai-colbert-small-v1 (96-d, via FastEmbed, ColBERT late interaction)
Review embeddings: BAAI/bge-small-en-v1.5 (384-d, via FastEmbed, per-finding)
API: FastAPI + Typer CLI
UI: Streamlit
Python: 3.11+

Benchmark results (baseline profile, 4 fixture products):

Mean latency: 2.2 ms / P95: 2.5 ms
Recall@3: 95.2% / Recall@5: 100%
MRR: 0.7204 / NDCG@5: 0.7920

Run it yourself:

git clone https://github.com/dvy246/qdrant-multivector.git
cd qdrant-multivector
uv sync --extra dev
docker compose up -d qdrant
EMBEDDING_BACKEND=deterministic uv run engine init-qdrant
EMBEDDING_BACKEND=deterministic uv run engine ingest --fixtures
EMBEDDING_BACKEND=deterministic uv run engine search \
  "Waterproof black hiking boots with good arch support" --user user_a

References

Qdrant vectors and multivectors documentation - the definitive reference for multivector configuration, MAX_SIM comparator behavior, and named vector field setup.
Qdrant hybrid queries and prefetch API - documents the query_points prefetch parameter used in the multi-stage search flow.
Qdrant payload indexing documentation - covers keyword, bool, and float index schema types used in create_payload_indexes.
Qdrant quantization documentation - covers scalar INT8 and binary quantization config, always_ram, and quantile settings.
FastEmbed ColBERT documentation - covers LateInteractionTextEmbedding, model selection, and token-matrix output format.
Hugging Face SigLIP model card - covers google/siglip-base-patch16-224, the vision and text tower architecture, and pooler output shape.
answerdotai/answerai-colbert-small-v1 - the specific FastEmbed ColBERT model used for 96-dimensional token-level text embeddings.
BAAI/bge-small-en-v1.5 - the BGE model used for 384-dimensional review finding embeddings.
Women's E-Commerce Clothing Reviews dataset - the real product dataset derived from Kaggle, used in the Streamlit demo.

Building Long-Running Claude Managed Agents: Why State Matters More Than Compute

Divy Yadav — Thu, 25 Jun 2026 10:50:30 +0000

A build story with real code, real failures, and the specific reasons one sandbox provider fixed problems I didn't know I had.

At 9:03am on a Tuesday, my research agent said hello and stared at an empty /workspace/.

Six hours of analysis from the night before. Gone.

The cloned repository. The installed packages. The notes it had spent hours writing. Gone.

I had assumed that if an agent stopped working for the night, it could simply continue the next morning. That was wrong.

Over the next three weeks, I rebuilt the same workflow on Tensorlake, Cloudflare, and Daytona to figure out what had happened. The hardest part of running Claude Managed Agents isn't the model. It's everything underneath it.

This is the exact code I ran, the things that broke, and the mistake that cost me two weeks to understand.

What Claude Managed Agents is, before anything else

If you've never built with Claude Managed Agents, the architecture needs a minute. Skip this if you already know it.

Anthropic runs the reasoning. You run the execution.

The agent loop, session state, work queue, and retry logic all live on Anthropic's infrastructure. You configure a Self-hosted Environment in the Claude Console. When your application starts a session, Anthropic queues the work, your orchestrator picks it up, spins up a sandbox, and the model starts issuing tool calls into that sandbox.

Every bash, read, write, grep, and edit call executes inside an environment you own. Anthropic never touches it. You decide what that environment looks like, what it can access, and what happens between sessions.

Anthropic's intelligence is fixed. Your engineering determines whether that intelligence has a stable, stateful environment to work in, or a clean slate that forgets everything the moment it goes idle.

What I was building and why it mattered

I needed an agent that could do real deep-work research on a codebase: clone a repository, read through the module structure, build an understanding of how the pieces fit together, write notes, and propose refactoring strategies. The kind of work that takes a senior engineer a full day and an AI agent about six hours.

The key constraint: the agent couldn't do this all at once. Sometimes I'd kick off a session at 8pm, let it run until midnight, and pick it back up the next morning. The filesystem it had built during that first session — the analysis notes, the installed tools, the half-read source files — had to be there when the next session started. Rebuilding from scratch each time wasn't viable.

That constraint is what drove every provider decision I made.

The requirements I didn't know I had

At the start I thought I needed a Linux environment that could run Claude Managed Agents. By the end, I realized I actually needed three things. I found them all in one place, but not until I had looked in two others first.

A filesystem that survived between work sessions.
Near-zero cost while the agent was idle.
The ability to branch from an already-completed analysis state.

I did not discover all three requirements on day one.I discovered them one mistake at a time.

How a session actually starts: the code before the sandbox

You drive a session through the reference orchestrator using a simple command:

make session PROMPT="Clone the repository at github.com/tensorlakeai/tensorlake. \
  Read through the module structure. Write a summary to /workspace/analysis.md. \
  Note any components that look like they could be simplified."

The orchestrator sends this prompt to Anthropic as a new session. Anthropic picks it up, starts the agent loop, and immediately begins issuing tool calls. Those tool calls arrive at your sandbox. The agent reads files, runs bash commands, writes notes. The session runs until the task is complete or you stop it.

The agent stream looks roughly like this as it runs:

[thinking] The repository appears to be a Python SDK for...
[bash] git clone https://github.com/tensorlakeai/tensorlake
[bash] ls -la /workspace/tensorlake/
[read] /workspace/tensorlake/tensorlake/sandbox.py
[write] /workspace/analysis.md
[thinking] The Sandbox class handles...

Each bracketed event is a tool call going into your sandbox. The session accumulates state inside /workspace/ across all those calls. By the end of a six-hour session, that directory contains the cloned repo, installed packages, analysis files, and intermediate notes. That's the state that needs to survive overnight.

Build 1: Cloudflare

My first assumption was that I needed a platform that could efficiently run Claude Managed Agents. Cloudflare is optimized for high-concurrency execution. My problem turned out to be different.

The agent I was building accumulated hours of filesystem state between bursts of work. Notes, cloned repositories, installed dependencies, and intermediate analysis all needed to survive overnight. Cloudflare's execution model wasn't designed around that requirement.That was the first time I realized I wasn't looking for compute.

I was looking for persistent state.

Build 2: Daytona

The second build solved part of the problem.The agent could accumulate state throughout a session, which initially felt like progress.

Then I wanted to test three different refactoring strategies starting from the same six-hour analysis. Instead of branching from that state, I found myself repeating the setup work each time: rebuilding context, reinstalling dependencies, and re-running analysis before I could begin the actual experiment.

That was when I discovered my second requirement.Preserving state wasn't enough.I also needed a way to branch from an existing state without repeating hours of work.

Build 3: Tensorlake

The first thing that caught my attention was not a feature.

It was an architectural decision.

Most platforms preserve state by keeping compute alive.

This one treated compute and state as separate problems.

The docs described a suspended sandbox that could preserve its state and resume in approximately 0.6 seconds. That was the first time I saw a design that directly addressed the problem I'd been running into.

I wanted to know whether it actually worked.

I started with the problem that had sent me down this path in the first place.Could an agent suspend overnight and resume with its state intact?

It could.

And once I tested checkpointing and branching, I finally had both things I'd been looking for.That's when the architecture started to make sense.

How the webhook architecture works

The deployment model here is different from a traditional always-on server, and understanding it made everything else click.

The orchestrator itself runs inside a Tensorlake sandbox with a public HTTPS endpoint. Anthropic pushes incoming work to that endpoint via webhook. When there's no traffic, the orchestrator sandbox suspends. When a new work item arrives, it wakes in under a second, processes the request, and creates a worker sandbox for that session.

Two independent lifecycles:

The orchestrator sandbox suspends when idle, preserving its memory state including the running uvicorn process. It doesn't accumulate per-session filesystem state, so suspending it between work items costs only storage rates.

The worker sandboxes — one per session — accumulate filesystem state throughout a session and suspend when the session ends. Their state is preserved in storage, not held by running compute.

Neither one has to stay alive to preserve the other's state. On every other platform I'd tried, "preserve state" meant "keep something running." Here it means "checkpoint and stop billing."

How session routing works

When you call make session PROMPT="...", the orchestrator sends a new session to Anthropic. Anthropic validates it, adds it to the work queue, and pushes a work item payload to your orchestrator's webhook endpoint.

That payload contains three things the orchestrator needs: the session ID, the work ID, and the environment ID. The orchestrator wakes, reads the payload, and creates a worker sandbox from your registered image:

from tensorlake.sandbox import Sandbox

sandbox = Sandbox.create(
    name=session_id,
    image="agent-cli",
    cpus=2.0,
    memory_mb=4096,
    timeout_secs=3600,
)

sandbox.start_process(
    "bash",
    ["-lc", "exec python3 /opt/sandbox_entrypoint.py > /tmp/runner.log 2>&1"],
    env={
        "ANTHROPIC_ENVIRONMENT_KEY": environment_key,
        "ANTHROPIC_SESSION_ID": session_id,
        "ANTHROPIC_WORK_ID": work_id,
        "ANTHROPIC_ENVIRONMENT_ID": environment_id,
    },
)

Two things tripped me up here. Credentials belong in start_process(env={...}), not Sandbox.create() — they're session specific, not image-level configuration. And sandbox names must be valid slugs. Neither was hard to fix once I knew what was happening, but both cost me time.

The worker sandbox runs sandbox_entrypoint.py, which attaches to the Anthropic session and begins consuming tool calls. From that point, the agent has a full Linux environment: bash, git, Python, and whatever else you built into the image.

Building the agent image

Every tool the agent needs has to live in the image before the session starts:

from tensorlake import Image

image = (
    Image(name="agent-cli", base_image="tensorlake/ubuntu-minimal")
    .run("apt-get update && apt-get install -y ca-certificates curl git gh python3 python3-pip")
    .run("pip install --break-system-packages 'anthropic>=0.103' 'httpx>=0.27'")
    .copy("sandbox_entrypoint.py", "/opt/sandbox_entrypoint.py")
    .workdir("/workspace")
)
image.build(registered_name="agent-cli")

I added my analysis tools here: Python packages, jq, ripgrep. If it needed to run inside the sandbox, it lived in the image.

How suspend and resume actually work

When a session ends, the worker sandbox suspends. Not terminates. The process state and filesystem are checkpointed to storage. Compute billing stops. The sandbox sits there as a stored snapshot until the next session begins.

Setting RESUME_SUSPENDED_SESSIONS=true tells the orchestrator what to do when the next work item arrives for an existing session: restore the suspended sandbox instead of creating a fresh one from the base image.

I set the flag and expected the sandbox to wake up immediately.

Nothing happened.

For a few minutes I thought the integration was broken. It wasn't. The flag doesn't trigger a resume directly. A new incoming webhook does. The flag just tells the orchestrator which action to take when that webhook arrives. Once I understood that, the behavior made complete sense .Why would a sandbox wake up before there's work to do?

The practical difference between fresh and resumed is simple: a fresh sandbox starts with a clean /workspace/. A resumed sandbox starts with whatever the last session left there.For a research agent that had already spent hours building an understanding of a codebase, those are completely different starting points.

The session ran overnight, suspended, and resumed the next morning with its state intact. The filesystem was exactly where I had left it.

Checkpointing and parallel exploration

Resume solved the overnight persistence problem. Checkpointing solved the one I'd discovered the hard way during the Daytona build.

After the full analysis was complete, I created a checkpoint and launched three sandboxes from it:

snap = sandbox.checkpoint()

children = [
    Sandbox.create(
        snapshot_id=snap.snapshot_id,
        name=f"{session_id}-strategy-{i}",
        cpus=2.0,
        memory_mb=4096,
    )
    for i in range(3)
]

Each child started from exactly the same state as the parent. The cloned repository was already there. The installed dependencies were already there. The analysis notes were already there. Nothing needed to be rebuilt.

Instead of repeating six hours of analysis three separate times, I ran three parallel experiments from the same verified baseline. By the next morning I had three implementations to compare.

The fork I'd dismissed as an edge case turned out to be the feature that mattered most.

Without it, every experiment required repeating hours of work. With it, experimentation became nearly free. This generalizes beyond refactoring: any time you want to test N variations with different prompts, different dependency versions, different approaches — you run the expensive setup once, checkpoint, and branch into as many parallel experiments as you need.

The cost of exploration drops to the cost of the experiments themselves.

Why I chose Tensorlake, stated plainly

My agent needed two things: zero idle compute cost while preserving filesystem state overnight, and the ability to fork mid-session to explore refactoring strategies in parallel. Those two requirements narrowed the field to one option.

I didn't choose Tensorlake because it was the fastest or the cheapest. I chose it because it directly addressed the two problems I was trying to solve.

The suspend/resume workflow preserved state without keeping compute running. Checkpointing made it possible to branch from an existing analysis instead of repeating hours of setup work. I tested both and they worked exactly the way I needed them to.

The hard part wasn't choosing a provider. It was understanding my requirements. Once I understood those, the decision became obvious.

How to Pick the Right Sandbox

	Tensorlake	Daytona	Cloudflare Containers	Cloudflare Isolates
Idle cost	Suspended at storage rate	VM stays alive	Scales to zero after sleep timeout	No filesystem to preserve
Filesystem on resume	State preserved	Full sandbox creation from snapshot	Disk is wiped on restart	No persistent filesystem
Resume speed	~0.6s memory restore	Full sandbox creation	Full VM restart	Milliseconds
Mid-session fork	Yes: checkpoint() + N	Yes, experimental fork method in core SDK (not documented in Claude integration guide)	Not documented	No
Full Linux	Yes	Yes	Yes	No
Self-hosted option	No	Yes	No	No
Concurrent scale	Up to 100k	Not published	Not published	Tens of thousands
Setup effort	Moderate	High	One-click	One-click

Does your agent accumulate state that's expensive to rebuild between bursts?

Tensorlake. This is the exact problem suspend/resume was designed to solve. If your agent starts fresh every session, continue to the next question.

Do you need to explore multiple approaches from the same mid-session state?

Tensorlake. Checkpointing and branching from an existing state eliminates hours of repeated setup work. If not, continue.

Do you need massive concurrency?

Cloudflare isolates are designed for very large-scale concurrent execution.

Do you need self-hosted infrastructure?

Daytona is the strongest fit if infrastructure ownership or data residency is a hard requirement.

How important is setup speed?

Cloudflare offers the fastest path to a prototype. Tensorlake requires more initial setup but provides capabilities that become valuable once agents accumulate long-lived state.

Conclusion

When I started, I thought I was choosing a sandbox provider.What I was actually discovering were my requirements.

First I learned I needed persistent state. Then I learned I needed a way to branch from that state without repeating hours of work.

Tensorlake was the first platform I tried that treated compute and state as separate problems.

Once I understood that distinction, the decision became obvious.The hardest part of running long-lived agents isn't getting them to work. It's making sure their work survives.

References

Tensorlake: Run Claude Managed Agents on Tensorlake Sandboxes — https://docs.tensorlake.ai/sandboxes/claude-managed-agents
Tensorlake: Sandbox Images (Image builder SDK) — https://docs.tensorlake.ai/sandboxes/images
Daytona: Run Claude Managed Agents on Daytona — https://www.daytona.io/docs/en/guides/claude/claude-managed-agents/
Cloudflare: Announcing Claude Managed Agents on Cloudflare — https://blog.cloudflare.com/claude-managed-agents/
Anthropic: Claude Managed Agents Self-Hosted Sandboxes — https://platform.claude.com/docs/en/managed-agents/self-hosted-sandboxes
Concurrency sandbox benchmarks results: https://platform.computesdk.com/scale-invitational

[Boost]

Divy Yadav — Tue, 16 Jun 2026 10:57:41 +0000

Divy Yadav

Jun 15

Why Most Multi-Agent AI Systems Waste 90% of Their Time (And How to Fix It)

#ai #programming #webdev #software

15 min read

Why Most Multi-Agent AI Systems Waste 90% of Their Time (And How to Fix It)

Divy Yadav — Mon, 15 Jun 2026 12:32:52 +0000

Most engineers treat multi-agent speed as a concurrency problem. It is not. The bottleneck is setup time, and memory snapshots change the math entirely.

Most engineers think multi-agent performance is a concurrency problem.

I did too.

So when five AI agents running in parallel barely outperformed a sequential run, I assumed something was wrong with my orchestration.

I was looking in the wrong place.

Each agent was spending more time preparing to work than actually working.

The fix wasn’t more threads, better async code, or a faster model.

It was a memory snapshot.

And once I saw where the time was really going, an entire class of multi-agent bottlenecks suddenly made sense.

Here is what that looks like, what took me three iterations to get right, and where it still has rough edges.

Let’s get the mental model first.

What This Does (30 Seconds)

The idea is straightforward: instead of five agents each spending 90 seconds installing the same tools, install them once, freeze that environment, and stamp out five identical copies.

Each copy runs a different analysis in parallel. A lead LLM reads all five results and tells you what to fix first.

In code:

Creates one Linux VM, installs code analysis tools (bandit, radon) and writes a sample Python project
Freezes the entire VM state into a memory snapshot (filesystem, memory, running processes included)
Forks 5 independent copies, each agent assigned a different analysis task (Security, Complexity, Docstrings, Tests, Structure)
Runs all 5 in parallel via asyncio.gather, finishing in seconds instead of minutes
Feeds all results to a lead LLM that produces a single prioritized fix list

Setup time is paid once, upfront, before any agent runs. The rest of this article explains how.

Why Sandboxes Matter for Agent Workloads

If you have not worked with sandboxes before: think of one as a disposable computer that lives in the cloud.

You spin it up, run whatever code you need inside it, and throw it away when you're done. It has its own filesystem, its own processes, its own network. Nothing it does can touch your machine or any other sandbox running at the same time.

In short: Sandboxes provide the agent with a secure and isolated enviornement

That isolation is the whole point. Your agent can install packages, write files, crash badly, or spin up a browser, and none of it bleeds out. When the task is done, you terminate the VM and it is gone.

The next agent starts clean.

Most agent frameworks treat the execution environment as an afterthought. The LLM call is the interesting part. The environment is just "wherever the code runs."

That works fine for single-turn tasks. It breaks down fast for anything multi-step.

When an agent needs to install packages, write intermediate files, maintain a browser session across multiple pages, or resume a task from a different machine, you need the execution environment to behave like a persistent object, not a function call that resets on every invocation.

Tensorlake gives each agent a MicroVM backed by Firecracker and CloudHypervisor, optimized for fast boot times and strong isolation. Each sandbox is a full Linux VM. It boots in hundreds of milliseconds, persists filesystem and memory state across sessions, and can be snapshotted at any point in its lifecycle.

Tensorlake also lets you spin up multiple sandboxes in parallel for concurrent agent execution, and honestly it is one of my favourite things about it.

it also ranks in the top 5 of SandboxBenchmarks.

What changes the math is a single question: what does the snapshot actually capture?

Two Kinds of Snapshots. Very Different Behavior.

Quick vocabulary before the details. Tensorlake sandboxes have four lifecycle modes.

An ephemeral sandbox runs a task and disappears when done, with no name and no persistence between runs.
A named sandbox outlives the process that created it and can be suspended then reconnected to from any machine. Suspend freezes the VM exactly as it is and resume brings it back to that same state.
A snapshot is that frozen moment saved as a reusable artifact.
A fork is a snapshot restored into a fresh, independent VM.

This project uses the last two.

Suspend and Snapshot both preserve state, but serve different purposes : Suspend is for pausing this sandbox to resume later, while a snapshot is a reusable artifact for retrying from a checkpoint or cloning an environment.

Tensorlake supports two checkpoint types. Most tutorials only mention one.

CheckpointType.FILESYSTEM captures disk state only. Restore from it and the new sandbox does a full cold boot: processes restart from scratch, packages get re-imported. Your pip installs survive. Nothing that was in memory does.
CheckpointType.MEMORY is different. It captures disk state, VM memory, and all running processes. The restored VM resumes mid-stride, exactly as the source was at checkpoint time. No boot sequence. No re-initialization. If Python had already imported bandit, the fork starts with it loaded. The environment is not rebuilt. It is copied.

The checkpoint type is not a performance detail. It determines whether your fork is a clone or a restart.

The default when you call sandbox.checkpoint() with no arguments is filesystem. That is the wrong choice for a parallel swarm where agents share a prepared environment. You want memory.

One more constraint worth knowing upfront: for memory snapshots, resources (CPUs, RAM) are baked into the snapshot at checkpoint time. You cannot override them when creating forks. Set the right cpus and memory_mb on the base sandbox before you checkpoint. Every fork inherits them automatically.

The Architecture

The pattern has five distinct phases. Each one has a single responsibility.

Phase 1 — Base Snapshot: Spins up a single baseline sandbox, installs analysis tools (bandit, radon), writes the target code, and checkpoints the entire running VM state using CheckpointType.MEMORY. The base sandbox is then terminated, leaving behind the reusable snapshot ID.

Phase 2 — Agent Forking: Restores 5 independent sandboxes concurrently from the base snapshot using sandbox.fork(...). Each fork is a warm start that inherits all installed tools, environment settings, and target files.

Phase 3 — Sequential Baseline (Timing): Runs each agent's analysis script (analyze.py) one-by-one inside its respective sandbox to measure sequential time as a benchmark denominator.

Phase 4 — Parallel Swarm: Executes all 5 agents concurrently using asyncio.gather(...). Each agent runs the same analysis script inside its isolated sandbox but with a different focus configuration passed via the PERSPECTIVE environment variable.

Phase 5 — LLM Aggregation: Collects the individual reports (Security, Complexity, Docstrings, Tests, Structure) alongside the timing data, and passes them to the lead LLM (GPT) to synthesize a single prioritized fix list.

Phase 1 runs once. Phases 2 through 4 run every time you want results. The fork is cheap. The base environment build is not, but you only pay that cost once per snapshot.

Phase 1: Build and Snapshot

The base sandbox installs the analysis tools, writes the target codebase into the VM, then snapshots the entire state. Every fork inherits both the tools and the target project automatically.

from tensorlake.sandbox import AsyncSandbox, CheckpointType

async def build_base_snapshot() -> str:
    async with await AsyncSandbox.create(
        name="base-swarm-env",
        cpus=2.0,
        memory_mb=2048,
        timeout_secs=600,
    ) as sandbox:

        # Install analysis tools. These are baked into the snapshot
        # and available to every forked agent at no extra install cost.
        result = await sandbox.run(
            "pip",
            ["install", "bandit", "radon", "--user", "--break-system-packages", "-q"],
            timeout=180,
        )
        if result.exit_code != 0:
            raise RuntimeError(f"pip install failed:\n{result.stderr}")

        # Write a sample Python project with intentional issues for agents to find.
        # All forks inherit this from the snapshot; no need to write per-agent.
        target_files = {
            "/workspace/target/auth.py": b'''
import subprocess
DB_PASSWORD = "hardcoded_secret_123"

def authenticate(user_input):
    return eval(user_input)

def run_command(cmd):
    return subprocess.call(cmd, shell=True)
''',
            "/workspace/target/logic.py": b'''
def classify(a, b, c, d, e, f, g, h):
    if a and b:
        if c or d:
            if not e and f:
                return "path_a"
            elif e and not f:
                return "path_b"
            elif g and h:
                return "path_c"
            else:
                return "path_d"
        elif g:
            return "path_e"
    return "path_f"
''',
        }
        for path, content in target_files.items():
            parent = "/".join(path.split("/")[:-1])
            await sandbox.run("mkdir", ["-p", parent])
            await sandbox.write_file(path, content)

        # Verify tools work before snapshotting.
        # A broken tool in the snapshot means broken forks.
        verify = await sandbox.run(
            "python3", ["-m", "bandit", "--version"]
        )
        if verify.exit_code != 0:
            raise RuntimeError(f"Tool verification failed:\n{verify.stderr}")

        snapshot = await sandbox.checkpoint(
            checkpoint_type=CheckpointType.MEMORY
        )

    # Context manager terminates the base sandbox here.
    if snapshot.status.value != "completed":
        raise RuntimeError(f"Snapshot failed: {snapshot.status.value}")

    return snapshot.snapshot_id

The async with pattern guarantees terminate() is called on exit, including on exceptions. Without it, any exception before a manual terminate() call leaves an orphaned VM running in the background. TensorLake's async documentation shows this pattern explicitly.

result.exit_code comes from CommandResult, the SDK's return type for run(). It has stdout: str, stderr: str, and exit_code: int. Note that stdout is already a string, not bytes, so no .decode() is needed anywhere.

The status check after checkpoint(): SnapshotStatus is an enum, so .value gives you "completed", "in_progress", or "failed". The documentation shows checkpoint() returns a SnapshotInfo with a status field. Checking that status before proceeding is a useful defensive practice. I learned this after a failed snapshot left me debugging downstream agent failures.

Phase 2: Fork and Run an Agent

This is the actual fork. The call is AsyncSandbox.create(snapshot_id=snapshot_id). No special fork() method. No copy-on-write API. Just create() with a snapshot ID. Every call produces a fully independent VM starting from that snapshot's frozen state.

PERSPECTIVES = ["Security", "Complexity", "Docstrings", "Tests", "Structure"]

async def run_agent(agent_id: int, snapshot_id: str) -> AgentReport:
    perspective = PERSPECTIVES[agent_id % len(PERSPECTIVES)]
    t_start = time.time()

    # cpus and memory_mb intentionally omitted.
    # For MEMORY snapshots, resources are inherited from the snapshot
    # and cannot be overridden at restore time.
    async with await AsyncSandbox.create(
        snapshot_id=snapshot_id,
        allow_internet_access=False,  # code analysis is offline; no outbound needed
        timeout_secs=120,
    ) as sandbox:

        await sandbox.write_file(
            "/workspace/analyze.py",
            ANALYSIS_SCRIPT.encode("utf-8")
        )

        result = await sandbox.run(
            "python3",
            ["/workspace/analyze.py"],
            env={"PERSPECTIVE": perspective},
            timeout=60,
        )

    elapsed = time.time() - t_start

    if result.exit_code != 0:
        raise RuntimeError(f"Agent {agent_id} failed:\n{result.stderr}")

    output = json.loads(result.stdout.strip())
    return AgentReport(
        agent_id=agent_id,
        perspective=perspective,
        score=output["score"],
        finding=output["finding"],
        execution_time_s=elapsed,
    )

allow_internet_access=False is safe here because bandit and radon analyze source code and do not make network calls. This parameter is not locked by MEMORY snapshots. TensorLake's networking documentation recommends disabling outbound internet access for untrusted code.

The dispatch script gets written fresh into each forked VM via sandbox.write_file(). Each agent's VM is fully isolated: writing to /workspace/analyze.py in fork 0 has no effect on fork 1. The target project files are already there, inherited from the snapshot.

Since result.stdout is already a Python string, json.loads(result.stdout.strip()) works directly. The .strip() handles the trailing newline from print() inside the sandbox.

Phase 3: Sequential First, Then Parallel

The sequential baseline exists for one reason: to give the speedup calculation a real denominator. Without it, you have a time with no context.

async def run_sequential(snapshot_id: str, count: int) -> SwarmResult:
    reports = []
    for i in range(count):
        reports.append(await run_agent(i, snapshot_id))
    return SwarmResult(mode="sequential", ...)

async def run_parallel(snapshot_id: str, count: int) -> SwarmResult:
    # asyncio.gather returns a list of results when awaited.
    reports = await asyncio.gather(
        *(run_agent(i, snapshot_id) for i in range(count))
    )
    reports.sort(key=lambda r: r.agent_id)
    return SwarmResult(mode="parallel", ...)

asyncio.gather is what TensorLake's async documentation recommends for concurrent sandbox fan-out. The ThreadPoolExecutor approach works too (the sync Sandbox API supports it), but if you are already in an async context, gather is cleaner.

Phase 4:What the Analysis Script Does

The dispatch script runs inside each forked sandbox. It reads the PERSPECTIVE environment variable, routes to the right analysis function, and prints one JSON line to stdout. All five analyses are fully offline, with no network calls needed.

# ANALYSIS_SCRIPT — runs INSIDE each forked sandbox
import json, os, subprocess, ast, pathlib, sys

PERSPECTIVE = os.environ["PERSPECTIVE"]
TARGET = "/workspace/target"

def run_security():
    """bandit: find hardcoded secrets, unsafe eval, shell injection."""
    r = subprocess.run(
        ["python3", "-m", "bandit", "-r", TARGET, "-f", "json", "-q"],
        capture_output=True, text=True
    )
    try:
        data = json.loads(r.stdout)
    except json.JSONDecodeError:
        return {"score": 0, "finding": "bandit parse error"}
    issues = data.get("results", [])
    high = [i for i in issues if i.get("issue_severity") == "HIGH"]
    return {
        "issues": len(issues), "high": len(high),
        "score": max(0, 100 - len(issues) * 10),
        "finding": high[0]["issue_text"] if high else ("Minor issues" if issues else "Clean"),
    }

def run_complexity():
    """radon: cyclomatic complexity per function."""
    r = subprocess.run(
        ["python3", "-m", "radon", "cc", TARGET, "-j"],
        capture_output=True, text=True
    )
    try:
        data = json.loads(r.stdout)
    except json.JSONDecodeError:
        return {"score": 0, "finding": "radon parse error"}
    blocks = [b for file_blocks in data.values() for b in file_blocks]
    complex_blocks = [b for b in blocks if b.get("complexity", 0) > 5]
    avg = sum(b["complexity"] for b in blocks) / len(blocks) if blocks else 0
    top = f"{complex_blocks[0]['name']} (cc={complex_blocks[0]['complexity']})" if complex_blocks else "All within threshold"
    return {
        "functions": len(blocks), "complex_count": len(complex_blocks),
        "avg_cc": round(avg, 2),
        "score": max(0, 100 - len(complex_blocks) * 15),
        "finding": top,
    }

def run_docstrings():
    """ast: count functions and classes that lack docstrings."""
    total, documented = 0, 0
    for path in pathlib.Path(TARGET).rglob("*.py"):
        tree = ast.parse(path.read_text())
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                total += 1
                if ast.get_docstring(node):
                    documented += 1
    pct = int(documented / total * 100) if total else 100
    return {"total": total, "documented": documented, "score": pct,
            "finding": f"{documented}/{total} documented ({pct}%)"}

def run_tests():
    """Count test files relative to source files."""
    all_py = list(pathlib.Path(TARGET).rglob("*.py"))
    test_files = [f for f in all_py if f.stem.startswith("test_") or f.stem.endswith("_test")]
    ratio = len(test_files) / len(all_py) * 100 if all_py else 0
    return {
        "source_files": len(all_py), "test_files": len(test_files),
        "score": min(100, int(ratio * 2)),
        "finding": f"{len(test_files)}/{len(all_py)} files are tests ({ratio:.0f}%)",
    }

def run_structure():
    """ast: count functions, classes, imports across the codebase."""
    stats = {"functions": 0, "classes": 0, "imports": 0, "files": 0}
    for path in pathlib.Path(TARGET).rglob("*.py"):
        stats["files"] += 1
        tree = ast.parse(path.read_text())
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):          stats["functions"] += 1
            elif isinstance(node, ast.ClassDef):           stats["classes"] += 1
            elif isinstance(node, (ast.Import, ast.ImportFrom)): stats["imports"] += 1
    fpr = stats["functions"] / stats["files"] if stats["files"] else 0
    return {**stats, "functions_per_file": round(fpr, 1),
            "score": min(100, int(fpr * 20)),
            "finding": f"{stats['functions']} functions across {stats['files']} files"}

dispatch = {
    "Security":   run_security,
    "Complexity": run_complexity,
    "Docstrings": run_docstrings,
    "Tests":      run_tests,
    "Structure":  run_structure,
}

fn = dispatch.get(PERSPECTIVE)

if fn is None:
    print(json.dumps({"error": f"Unknown perspective: {PERSPECTIVE}"}))
    sys.exit(1)

result = fn()
result["perspective"] = PERSPECTIVE
print(json.dumps(result))

Two things worth keeping when you adapt this.

Parameters via environment variables: sandbox.run(env={"KEY": "val"}) passes per-command variables and avoids shell escaping issues when values contain spaces or special characters. It also keeps the dispatch script stateless, with no hardcoded perspective names inside the script itself.

JSON to stdout: the orchestrator reads result.stdout.strip() and passes it directly to json.loads(). The script has one job: print exactly one valid JSON line. Any other stdout output (debug prints, progress bars) breaks the parse. Keep it strict.

Phase 5: Lead Agent Synthesis

After all five agents return, a single GPT-4o call synthesizes their findings into a prioritized action list.

def aggregate_with_llm(parallel: SwarmResult, sequential: SwarmResult) -> str:
    client = OpenAI()
    speedup = sequential.total_time_s / parallel.total_time_s

    reports_block = "\n".join(
        f"[{r.perspective}] Score: {r.score}/100 | {r.finding}"
        for r in parallel.reports
    )

    prompt = (
        "You are a senior engineering lead reviewing a parallel code analysis report.\n\n"
        f"Agent Findings:\n{reports_block}\n\n"
        "Benchmark:\n"
        f"  Sequential : {sequential.total_time_s:.2f}s\n"
        f"  Parallel   : {parallel.total_time_s:.2f}s\n"
        f"  Speedup    : {speedup:.2f}x\n\n"
        "Provide: overall codebase health score, top three issues to fix immediately "
        "(with file and severity), recommended next actions, and one sentence on what "
        "the parallel speedup means for running this at scale."
    )

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content

The lead agent sees both the analysis findings and the timing benchmark in the same context. That is the reduce step in a map-reduce agent pattern: give the aggregator everything the workers produced, not just the domain data. The call is synchronous because there is nothing left to concurrently await at this point.

Where the Time Actually Goes

Both timelines contain the same agents doing the same work. What changes is when setup happens. These numbers are structural projections based on typical pip install times and sandbox warm-restore behavior, not measured results. Your numbers will vary by workload and network conditions. Run the demo to measure your case.

Without memory snapshots:

Agent 0: [setup ~90s][work ~8s]
Agent 1: [setup ~90s][work ~9s]
Agent 2: [setup ~90s][work ~8s]
Agent 3: [setup ~90s][work ~9s]
Agent 4: [setup ~90s][work ~8s]

Sequential total: ~490s
Parallel total:   ~100s  (setup still paid by each fork separately)

With memory snapshots (MEMORY type):

Base build:  [setup ~90s][checkpoint ~3s]  ← paid once, outside the loop
Agent 0: [warm fork ~1s][work ~8s]
Agent 1: [warm fork ~1s][work ~9s]
Agent 2: [warm fork ~1s][work ~8s]
Agent 3: [warm fork ~1s][work ~9s]
Agent 4: [warm fork ~1s][work ~8s]

Sequential total: ~48s
Parallel total:   ~10s

The speedup ratio looks similar on paper. The absolute time is not. At five agents the gap is 450 seconds versus 5 seconds of overhead. At fifty agents it is 4,500 seconds versus 50 seconds.

Setup time does not scale down with parallelism. It multiplies. The snapshot moves it outside the loop entirely.

The benchmark captures four numbers: sequential total time (the denominator), parallel total time (wall-clock from first fork to last return), speedup (sequential divided by parallel), and efficiency (speedup divided by agent count, multiplied by 100).

Efficiency is the one most benchmarks skip. A 4.2x speedup across five agents is 84% parallel efficiency: 16% is lost to fork startup, scheduling, and I/O contention. That number matters when you scale from five agents to fifty.

What the Code Does Not Handle

The demo covers the happy path. Three things to add before production:

LLM rate limits. Twenty or thirty concurrent agents all hitting the OpenAI API will trigger rate limit errors. The demo has no retry logic. Add exponential backoff before you scale.
Snapshot storage. Snapshots may incur charges depending on your plan. Use Sandbox.delete_snapshot(snapshot_id) when done. The demo has a CLEANUP_SNAPSHOT_ON_EXIT flag at the top of the file.
Agent error isolation. If one run_agent() coroutine raises inside asyncio.gather, the whole batch fails. In production, wrap each coroutine with asyncio.create_task() and handle errors per-agent.

When to Use This Pattern (And When Not To)

Use it when:

Multiple agents need the same environment
Their tasks are independent (no inter-agent communication mid-run)
Setup time is a meaningful fraction of total runtime
Reproducibility matters: every fork starts from an identical state

Skip it when:

Agents need to share state during execution. Forks are fully isolated. If agent 2 needs to react to what agent 1 found, use shared storage or message queues instead.
The task is fast enough for a single agent. Forking five sandboxes for a 3-second job adds overhead, not speed.
Environment setup takes under 5 seconds. The snapshot overhead only pays off when setup is the actual bottleneck.

Your situation	Right choice
Multiple agents, shared dependencies, independent outputs	Memory snapshot, fork N copies
Single agent, long task, needs to pause and resume	Named sandbox with suspend/resume
Pure browser automation, no code execution	Stagehand or BrowserBase
Stateless task, resets every run	Ephemeral sandbox, no snapshot needed
Environment setup under 5 seconds	Filesystem snapshot or skip snapshots

On filesystem performance: Tensorlake publishes performance benchmarks on their GitHub comparing sandbox execution times across providers. Refer to their repository for current numbers.

Running This

pip install tensorlake openai
export TENSORLAKE_API_KEY="your-key"
export OPENAI_API_KEY="your-key"
python3 agent.py

Free tier at cloud.tensorlake.ai, no credit card required. The demo takes 3-5 minutes end to end. After it runs, benchmark_results.json has the full per-agent timing data.

Phase 1 (base build and snapshot) runs once. If you want to run the benchmark multiple times, pass your existing snapshot ID directly and skip Phase 1. The snapshot persists between runs until you delete it.

What Actually Took Three Iterations

The first version had plain await sandbox.terminate() at the end of each function. Two exceptions during testing left sandboxes running and billing for idle compute. Switched to async with await AsyncSandbox.create(...) as sandbox: and that stopped.

The second version called sandbox.checkpoint(sandbox.sandbox_id). I had copied the pattern from a CLI reference (tl sbx checkpoint <sandbox-id>) and assumed the Python SDK matched. It does not. The Python instance method takes no positional arguments: sandbox.checkpoint(checkpoint_type=CheckpointType.MEMORY). That is it.

The third version was the first one that ran end to end, but with CheckpointType.FILESYSTEM by default because I had not read the snapshots documentation carefully. The benchmark looked reasonable. The forks were doing full cold boots and I was measuring them alongside the actual work. Switching to CheckpointType.MEMORY was the change that made setup time disappear from per-fork timing.

Small mistakes individually. What they share: Tensorlake's API is well documented, but the snapshot docs, the SDK reference, and the async docs are three separate pages. Read only the quickstart and you miss two of the three things that matter most for this pattern.

You can also check the complete project on my github here:

click_here

The Thing That Changes

Running the same five agents sequentially and then in parallel is one of those moments where the architecture becomes legible in a way that documentation does not fully convey.

The snapshot moves setup cost from inside the loop to outside it. The agents still do the same work on the same hardware. The savings come from not rebuilding an environment five times when it only needed to be built once.

Most multi-agent optimization advice focuses on LLM calls: batching, caching, cheaper models. That advice is right. But if you have five agents each spending 90 seconds on pip installs before making a single inference call, no amount of LLM optimization helps until you address setup time first.

The bottleneck was never the agents. It was rebuilding the same environment on every run. Snapshot it once, fork cheaply, and parallel execution finally delivers what you expected when you first wrote asyncio.gather.

References:

MCP Is Dead. The Downloads Just Don't Know It Yet.

Divy Yadav — Fri, 05 Jun 2026 12:54:22 +0000

30 CVEs in 60 days, a maintenance tax nobody warned you about, and what engineers are quietly switching to.

Your AI agent ran a query on a fake database last month.

It got real results. The tool worked perfectly. Your SSH keys left in the background.

The agent didn't flag it. The registry didn't catch it. Nobody warned you.

That's not a hypothetical. That's MCP in 2026, with 97 million monthly downloads and a Linux Foundation home.

The hype was real. So are the cracks.

First: what is MCP, and why should you care

If you've never built AI agents before, this matters. Skip it if you have.

Say you're building an AI assistant that needs to do real work:

Look up customer records in a database
Create tickets in Jira
Send a Slack message
Pull a file from Google Drive

Each of those lives in a different system. Different API, different auth, different data format.

To connect your AI to all of them, you'd write a custom integration for each one. Fine for two tools. Painful for ten. Then you switch models and rewrite everything.

This is the N×M problem: N tools multiplied by M AI models equals a mountain of glue code nobody wants to maintain.

MCP — the Model Context Protocol — solves that. Released by Anthropic in November 2024, it's an open standard that gives AI models one universal way to talk to external tools. You build an MCP server once around a tool, and any MCP-compatible AI can use it.

Your agent  →  MCP Client  →  MCP Server  →  Real Tool (Slack, Postgres, GitHub)

Three pieces:

MCP Host: your app (Claude Desktop, VS Code, a custom agent)
MCP Client: the component inside your app that speaks MCP, discovers tools, calls them
MCP Server: a small process wrapping a real tool, exposing it in a format any MCP client can use

That's it. The N×M problem disappears. One integration per tool, works with every AI.

The pitch was real. The adoption proved it.

How MCP went from zero to everywhere in 14 months

The adoption happened fast. Unusually fast.

Nov 2024 — Anthropic launches MCP. ~2M monthly SDK downloads.
Apr 2025 — OpenAI adopts it. Downloads jump to 22M.
Jul 2025 — Microsoft integrates it into Copilot Studio. 45M.
Nov 2025 — AWS adds support. 68M.
Mar 2026 — Every major AI vendor on board. 97M downloads. 10,000+ public MCP servers.

In December 2025, Anthropic donated MCP to the Agentic AI Foundation under the Linux Foundation, co-founded with OpenAI and Block. It stopped being Anthropic's protocol and became the industry's.

The "USB-C for AI" comparison spread everywhere. Everyone plugged in.

And then engineers started running it in production.

Why developers hate MCP

The complaints have been building in forums, GitHub issues, and private Slack channels for months. Not from people who misunderstood MCP. From people who ran it in production and got surprised by the same things.

"It's unauthenticated by default."

Out of the box, an MCP server trusts whatever connects to it.

No built-in check that the server is who it claims. No built-in check that the client is who it claims. You're responsible for adding that layer.

Most tutorials don't mention this.

"The STDIO transport executes arbitrary OS commands."

The official MCP STDIO transport runs any OS command you point at it to launch a server. Even when the server startup fails. No sanitization warnings. Nothing in the developer toolchain flags it.

OX Security documented this in April 2026.

Anthropic's response: expected behavior, sanitization is the developer's responsibility. LangChain said the same. Microsoft said the same.

Three major vendors. Same answer: your problem.

"The spec moves, community servers don't."
MCP servers published in community registries frequently fall behind spec updates. A server that worked last month may behave differently after the protocol updates. The registry has no enforcement mechanism. You find out in production.

"Every REST API I already have needs a new wrapper process."

Adding MCP to a tool that already has a clean REST API means building an entire MCP server around it.

That server needs to be:

Deployed and monitored
Updated when the underlying API changes
Secured separately from both the agent and the tool

For ten existing APIs, that's ten new processes to own. Month three of production, you feel every one of them.

"The registries are basically npm circa 2015."

In early 2026, OX Security cloned mcp-server-postgres and named it mcp-server-postgress (extra 's'). Functionally identical. Same queries, same responses.

Hidden inside: a payload that silently pulled SSH keys and environment files to an outside server.

They submitted it to eleven major MCP registries.

Nine published it. No automated security review. No source code analysis. Nothing.

"It eats my context window before the user says anything."

When your MCP client connects to a server, it loads the full tool schema into your context window — names, descriptions, parameters for every tool.

One server, five tools: ~500 tokens gone before the first message
Ten servers: 2,000–3,000 tokens gone before the first message

The model is already reasoning over a smaller budget. Before your user typed a word.

These aren't edge cases. They're the standard experience for anyone who's moved past a local demo.

The 3 problems that actually break production systems

The trust problem

Your agent has no way to verify an MCP server is who it claims to be.

The OX Security incident made this real: nine out of eleven registries accepted a typosquatted credential-stealing package. The malicious server functioned correctly. Ran database queries. Returned results. And silently pulled your SSH keys in the background. Nothing in the protocol flagged it.

Since January 2026, researchers have filed 30 CVEs against the MCP ecosystem in 60 days. Prompt injection through tool descriptions. Credential theft via config file reads. "Tool poisoning," where a server description manipulates the agent's next decision. These aren't exotic attack vectors.

Your agent can't tell your Postgres server from an attacker's. That's not a code bug. It's a design gap.

MCP was built for a trusted local environment. Production isn't that.

The wrapper tax

Every tool you connect to MCP needs its own MCP server. Ten tools means ten additional processes to own.

Each one needs to:

Stay in sync when the underlying tool's API changes
Be monitored for failures in production
Be secured separately from the agent and the tool itself
Be deployed as part of your infrastructure

For the first two tools, manageable. Month three with fifteen integrations, it's a job.

The N×M problem is solved. The "N new processes" problem quietly replaced it.

The context window bill

Tool schemas are not free. They're tokens. And they arrive before your user's message does.

A team building a customer service agent connected to ten MCP servers found their available reasoning budget had shrunk by 30% before the first user question arrived. Same model. Same prompts. Just more tools.

In a long multi-step agent session, schema tokens compound. Quality drifts. Costs climb. Most teams don't trace this back to tool schema overhead until they look at what's actually in the context window.

What engineers are using instead

A few patterns have emerged for teams that ran into the problems above.

Direct REST API calls

For tools with a clean existing API, skip MCP entirely. Call the API directly from your agent. No new server to maintain, no schema overhead, existing auth covers it.

Works well when you control the tool and the API is stable. Doesn't scale when you need multiple AI systems to share the same integrations.

Native provider tool use

Anthropic and OpenAI both have built-in tool calling that needs no MCP infrastructure. You define the tool schema inline, pass it with the request, the model calls it.

No server process
No registry
Your auth sits directly on the call

Most teams running focused single-purpose agents in 2026 are doing this. Simpler to reason about, harder to share across systems.

UTCP (Universal Tool Calling Protocol)

UTCP skips the wrapper entirely. Instead of wrapping a tool in an MCP server, it calls the tool's existing HTTP endpoints directly, with a discovery layer on top.

As of early 2026:

1,000+ GitHub stars
Implementations in Python, Go, and TypeScript
Growing community from teams that wanted lower latency and less infrastructure overhead

Best for teams with well-designed existing APIs who don't want a separate server layer. Not a full MCP replacement if you need the ecosystem breadth — but for many production use cases, materially simpler.

MCP with a gateway layer

For teams committed to MCP, the answer to most of the problems above is an MCP gateway — a controlled layer between your agent and your servers.

Your agent  →  MCP Gateway  →  MCP Server 1  →  Tool
                            →  MCP Server 2  →  Tool
                            →  MCP Server N  →  Tool

A gateway handles:

Authentication — verifies server identity before your agent calls anything
Tool filtering — loads only schemas relevant to the current task, not all of them
Audit logging — records every tool call for compliance and debugging
Rate limiting — stops runaway tool calls from blowing your budget

As of April 2026, 86–89% of AI agent pilots fail before reaching production. Governance gaps and audit visibility are the two most common reasons. A gateway is what closes both.

So do we actually need MCP?

Yes. With caveats that matter.

Use MCP when:

Multiple AI systems need to share the same tools
You're a SaaS company giving AI agents access to your product
You need dynamic tool discovery across a large integration ecosystem

Skip MCP when:

You're building a focused agent with two or three tools you already own
Your tools have clean REST APIs you control
You need low latency and minimal infrastructure overhead

The "MCP for everything" era is over. It's the right call when standardization pays off at scale. When you just need your agent to hit an API you already control, MCP is overhead pretending to be infrastructure.

Cheat sheet: what to actually do

Your situation	What makes sense
Local dev, one or two tools, just exploring	Bare MCP or native tool calls. Don't over-engineer.
Agent using tools you own, clean REST APIs	Direct API calls or native tool use. Skip MCP overhead.
Production agent, 5+ tools, or external users	MCP with a gateway. Authentication is not optional.
Enterprise, compliance, or regulated industry	MCP gateway with audit logs and SSO. Non-negotiable.
Pulling from community MCP registries	Treat every server as untrusted. Verify before deploying.
Need to share tools across multiple AI systems	MCP is the right call. This is exactly what it's for.

The actual state of things

MCP isn't going away. The downloads are real. The Linux Foundation governance is serious. Multi-vendor adoption means the protocol has institutional staying power.

But the MCP of early tutorials — install a community server, plug it in, done — that version is dead.

It was never safe for production. It was never meant to be.

The engineers moving to UTCP or direct API calls aren't abandoning MCP because it failed. They're routing around the parts that weren't built for what they're building.

I keep coming back to the OX Security test. Nine out of eleven registries. No automated review. The agent called the fake server, ran its queries, and handed over credentials it didn't know it was handing over.

Your agent does what it's told by the tools it trusts.

MCP hasn't fully answered how it decides what to trust. Until it does, treat every community MCP server the way you'd treat a random npm package in 2015.

You know how that era ended.

I Built a Stateful Research Agent Inside a Sandbox. Here's What the Numbers Actually Looked Like.

Divy Yadav — Wed, 27 May 2026 04:37:05 +0000

Three steps into a multi-page research task, the agent lost everything.

Not a crash. Not a thrown exception.

The function returned, context reset, and the pricing data it had just collected vanished.

This failure is predictable: stateless execution environments were never built to hold state across browser sessions that run for twenty minutes.

You hit it eventually, usually at the worst moment.

The two standard workarounds are both annoying. Stuffing state into the prompt works until token costs starts becoming an issue. An external state store solves the problem but now you are maintaining another service.

I had been using E2B for short-lived code execution. It handles that well, and they have added persistence features over time, including early-stage snapshot support. But for agents that need to pause mid-task and resume from a different process, state management is still mostly on you.

Someone in my Discord mentioned TensorLake. I opened the docs and decided to build against this specific problem.

In this article, I will walk you through the steps using which you can build a desktop using an agent in a sandbox.

Let's start with setting up.

Visual Explanation First

Setup

What caught my attention first: named sandboxes with suspend() and resume() that preserve the full VM state, not just files, but running processes and open browser sessions. Sub-second resume, according to their docs.

Ten minutes from zero to running:

pip install tensorlake
tl login   # or TENSORLAKE_API_KEY env var

Free tier, no credit card.

from tensorlake.sandbox import Sandbox

sandbox = Sandbox.create(
    name="research-agent",
    cpus=2.0,
    memory_mb=4096,
    secret_names=["OPENAI_API_KEY"],
    image="tensorlake/ubuntu-vnc",
)

The tensorlake/ubuntu-vnc image is what gives you a real desktop and Firefox inside the VM. You need an actual browser because modern pricing pages heavily use client-side rendering and bot detection that stops headless scrapers cold. Firefox inside a sandbox just looks like a person browsing.

Important: Playwright is not pre-installed in ubuntu-vnc. Install it before the agent runs:

sandbox.run("pip", ["install", "playwright"])
sandbox.run("playwright", ["install", "chromium"])

Two to three minutes on first setup. After that, packages persist across suspend/resume so you pay the cost once.

Latency: What I Actually Measured

First sandbox was running in roughly 800-900ms from the Sandbox.create() call to status running.

Here is where time actually goes:

Sandbox creation:        ~800ms          (named sandbox, first time)
Sandbox resume:          ~400ms          (from suspended state)
LLM call (GPT-4o):       2,000-4,000ms   (per step, dominates everything)
Browser screenshot:      ~300ms          (capture + transfer)
Page load in sandbox:    1,000-2,000ms   (varies by site)
File read/write:         <50ms           (block-based storage)
Sandbox suspend:         ~200ms

The LLM calls dominate by a large margin. Sandbox overhead is not the bottleneck. The main optimization is batching browser operations before each model call rather than interleaving individual round trips.

Tensorlake publishes a SQLite filesystem benchmark claiming 1.6-1.9x faster I/O than E2B and Modal. Self-reported numbers. I could not independently verify them. What I can say is that the block-based storage felt responsive for frequent small writes, which is exactly the pattern a research agent uses when checkpointing after every step.

Computer Use: What Worked and What Didn't

The desktop API itself is clean:

with sandbox.connect_desktop(password="tensorlake") as desktop:
    png_bytes = desktop.screenshot()
    desktop.move_mouse(640, 400)
    desktop.click()
    desktop.type_text("pinecone.io")
    desktop.press("Return")

Screenshot as PNG bytes, decode it, figure out where to click, send coordinates. Each browser interaction takes 1-3 seconds depending on page load. Slow compared to an API call. But it works on pages that block scrapers, because from the server's side it is just a person using Firefox.

The problem: coordinates assume a fixed layout, and layouts do not stay fixed.

Weaviate's pricing page ran an A/B test between two of my agent's steps. The toggle moved 30px down. The agent clicked empty space. No error, no exception. Just a screenshot showing nothing happened, and twenty minutes of debugging before I identified the offset.

The fix: pass screenshots to GPT-4o Vision to identify element positions dynamically rather than hardcoding coordinates. Adds about 2 seconds per interaction, handles layout drift reliably. Worth it for reliability; too slow for high-frequency operations.

When the DOM is accessible, Playwright inside the sandbox is the better path:

result = sandbox.run(
    "python",
    ["-c", """
import asyncio
from playwright.async_api import async_playwright

async def get_pricing():
    async with async_playwright() as p:
        browser = await p.chromium.launch()
        page = await browser.new_page()
        await page.goto("https://pinecone.io/pricing")
        pricing_text = await page.inner_text(".pricing-section")
        print(pricing_text)
        await browser.close()

asyncio.run(get_pricing())
"""]
)

The hybrid strategy I landed on:

Situation	Approach	Why
Site with bot detection	Vision + coordinates	Playwright gets blocked
Accessible DOM	Playwright directly	Faster, no coordinate drift
Unknown or variable layout	Screenshot + GPT-4o Vision	Resolves position dynamically
High-frequency operations	Playwright only	Vision adds ~2s per call

Use vision as a fallback, not a first tool. Vision handles layout variation. Playwright handles speed. Neither does both well.

Statefulness: The Part That Actually Mattered

After three steps (Pinecone free tier limits noted, $70/mo Starter plan recorded, Weaviate docs started), I called sandbox.suspend().

The sandbox froze. Filesystem, memory, running browser: all paused. Twelve minutes later, from a different terminal:

sandbox = Sandbox.connect("research-agent")
sandbox.resume()

About 400ms. The Weaviate pricing tab was still open. Tensorlake's suspend/resume preserves the full VM state, including memory and running processes.

Everything written to /workspace/research_notes.json was intact.

The workflow I settled on: write state explicitly after each meaningful step, then suspend.

# After each step, before suspending:
sandbox.write_file(
    "/workspace/state.json",
    json.dumps({
        "pinecone_pricing": pinecone_data,
        "weaviate_started": True,
        "next_url": "https://weaviate.io/pricing"
    }).encode()
)
sandbox.suspend()

# On next invocation, from any process:
sandbox = Sandbox.connect("research-agent")
sandbox.resume()
state = json.loads(bytes(sandbox.read_file("/workspace/state.json")))
# picks up from state["next_url"]

The state file is the continuity mechanism. Not elegant, but it removes the need for an external database and the filesystem is fast, durable across suspend, and readable from any reconnecting process.

Scaling and Failure Handling

Sandbox.create() is a blocking synchronous call. For parallel workloads, wrap in concurrent.futures:

from tensorlake.sandbox import Sandbox
from concurrent.futures import ThreadPoolExecutor

def research_competitor(name, url):
    sandbox = Sandbox.create(
        name=f"research-{name}",
        cpus=1.0,
        memory_mb=2048,
        secret_names=["OPENAI_API_KEY"],
        image="ubuntu-vnc",
    )
    # ... agent logic ...
    result = sandbox.read_file("/workspace/report.json")
    sandbox.terminate()
    return result

competitors = [
    ("pinecone", "pinecone.io/pricing"),
    ("weaviate", "weaviate.io/pricing"),
    ("qdrant", "qdrant.tech/pricing"),
]

with ThreadPoolExecutor(max_workers=5) as executor:
    reports = list(executor.map(lambda c: research_competitor(*c), competitors))

Three concurrent sandboxes ran without delay. I have not tested at twenty or fifty. Their docs mention hundreds per second. Take that at face value until you have load data.

Note: Tensorlake's Python SDK v0.5.8 introduced native async APIs that offer a cleaner alternative to threading for I/O-bound orchestration. If you are on v0.5.8 or later, those are worth reaching for before wrapping synchronous calls in a thread pool.

Patterns worth building from day one:

Idempotent state writes. Write state after each meaningful step. If the agent fails mid-run, the next invocation reads the file and skips completed work. This does not happen automatically.

Checkpoint before risky operations. sandbox.checkpoint() creates a restorable snapshot. By default, snapshots preserve the filesystem state. Preserving full memory state is supported as an explicit option. Either way, you can restore into a fresh sandbox if an operation goes wrong:

# Filesystem snapshot (default)
snapshot = sandbox.checkpoint()

try:
    agent.navigate_to_pricing_page()
except Exception:
    # Restore filesystem state into a new sandbox
    sandbox = Sandbox.create(snapshot_id=snapshot.snapshot_id)

Named sandboxes. If the orchestration process dies, any other process reconnects with Sandbox.connect("sandbox-name") and resumes from the last written state.

Architectural boundary: Tensorlake provides the execution environment and runtime for agents: the VM, the filesystem, the process lifecycle, the networking. It is not an agent framework. Retry logic, circuit breakers, and LLM rate-limit backoff belong in the orchestration layer above it: LangChain, LlamaIndex, a custom harness, or whatever you are using to drive the agent. That separation is deliberate, not a gap.

The Mental Model

The part that shifted how I thought about the design:

┌─────────────────────────────────────────────┐
│                 Your Agent                   │
│    (LLM + tool calling logic)                │
└──────────────────┬──────────────────────────┘
                   │ tool calls
┌──────────────────▼──────────────────────────┐
│           Tensorlake Sandbox                 │
│  ┌──────────────────────────────────────┐   │
│  │ State Layer: /workspace filesystem   │   │
│  │  state.json, research_notes.json     │   │
│  └──────────────────────────────────────┘   │
│  ┌──────────────────────────────────────┐   │
│  │ Execution Layer: processes, scripts  │   │
│  └──────────────────────────────────────┘   │
│  ┌──────────────────────────────────────┐   │
│  │ Computer Use: VNC, screenshots, mouse│   │
│  └──────────────────────────────────────┘   │
└─────────────────────────────────────────────┘

The sandbox is not the agent. It is the stable environment the agent operates in. When it resumes, the environment is exactly where the agent left it. The agent's logic lives outside and reconnects to a world that did not reset.

That changes what you can build. An agent that runs for an hour, navigates fifteen pages, and writes a structured report is feasible when the execution environment outlasts the orchestration session. With purely ephemeral execution, it is not.

How It Compares

vs E2B:

Both use Firecracker microVMs. E2B markets sub-200ms cold starts; community reports put real-world p50 closer to 400-600ms. Tensorlake named sandbox creation was ~800ms in my testing.
E2B has added snapshot and pause-resume in recent releases. The statefulness gap is narrower than a year ago. Tensorlake's suspend/resume preserves the full running VM state, including open processes, browser sessions, all in under a second. E2B's memory snapshot support is still described as early-stage.
Tensorlake claims 1.6-1.9x faster filesystem I/O on their own benchmarks. Self-reported. For an independent reference: Tensorlake recently ranked top 2 across all three categories in the ComputeSDK sandbox benchmarks.
Neither provides DOM-level element selection at the SDK layer.

vs Modal:

Modal uses gVisor rather than Firecracker, designed around stateless function execution. Stateful long-running agents work but need more setup. Cold starts are around 1-1.5 seconds per their docs.

vs Stagehand (BrowserBase):

Stagehand has DOM-level selectors (CSS, XPath, natural language) via locator(). For pure browser automation, this is a real ergonomic advantage.
Tensorlake gives you a full VM. Code execution, file management, package installs, and browser use in the same environment. If that combination is what you need, the full VM model is worth the coordinate complexity.
Browser automation only? Stagehand is the more focused tool.

from tensorlake.sandbox import SandboxClient

client = SandboxClient()

for sb in client.list():
    print(sb.sandbox_id, sb.status)

What the Build Produced

By the end of the session, the agent had produced the comparison: Pinecone versus Weaviate pricing, extracted across seven pages, with notes preserved across two suspensions and a full restart of the orchestrating machine.

report_bytes = sandbox.read_file("/workspace/comparison_report.md")
print(bytes(report_bytes).decode("utf-8"))

Accurate. Correct tier names and numbers.

Tensorlake did not solve the hard parts: the retrieval logic, state schema, hybrid browser strategy. It stayed out of the way while those got built. Most of the infrastructure friction came down to state management, and most of that went away once the sandbox filesystem became the state store.

Three Things to Know Before You Start

Speed is a systems problem, not a sandbox problem. LLM calls account for the bulk of per-step latency. Optimize by batching browser operations before each model call, not by chasing sandbox startup time.

Design for interruption from day one. Write state after every meaningful step. Not because the sandbox will crash, but because resuming from a different process after an unexpected interruption is a real scenario, not an edge case.

Computer use is a primitive. The coordinate-based API works, but layout drift will break hardcoded positions. Use Playwright when the DOM is accessible. Fall back to vision when you need a real browser session. Do not automate full workflows with raw coordinates.

Is the sandbox infrastructure production-ready? Yes. Suspend/resume held up, filesystem persistence was consistent, and Firecracker isolation did what it was supposed to.

Is the computer use layer production-ready? Not without additional engineering. The raw coordinate API is a reasonable primitive, but element resolution needs to be built on top of it. A vision-backed click_element() in the SDK would change the story significantly. Until then, budget the time to build that layer yourself.

Worth using? Yes, if you go in with clear expectations about what the platform handles and what it leaves to you. That boundary is sharper than most, which makes it easier to work with once you have internalized it.

You can also check the complete project on my github here:

click_here

References

Tensorlake. Tensorlake Documentation & Sandbox SDK. https://tensorlake.ai

E2B. E2B Sandbox Infrastructure. https://e2b.dev

Modal. Modal Serverless Infrastructure. https://modal.com

Stagehand (BrowserBase). Stagehand Browser Automation. https://browserbase.com/stagehand

Amazon Web Services. Firecracker MicroVMs. https://firecracker-microvm.github.io/

Microsoft. Playwright Browser Automation. https://playwright.dev

Benchmark: https://www.computesdk.com/benchmarks/sandboxes/

LLMs, RAG, Agents, MCP: The AI Evolution You Actually Need to Understand

Divy Yadav — Wed, 20 May 2026 11:03:00 +0000

Most people still think AI is just a chatbot.

That idea is already outdated.

Modern AI systems browse the web, remember your preferences, execute code, query databases, call APIs, and coordinate workflows. They operate more like software employees than like a search bar.

This did not happen because models got smarter. It happened because the architecture changed.

Every layer of the modern AI stack exists because the previous layer had a real failure. Understanding what failed and why something new was built is the fastest way to understand how any serious AI product works today.

That is what this article covers. Every stage: LLMs, RAG, Agents, and MCP.

In this article, you will get a good idea about the full AI Evolution, and it took me a lot of research and work for this one.

So if you want more such information about AI, consider subscribing to my newsletter, where you will get noise-free information every week

Link for the newsletter: Newsletter

Stage 1: The LLM Era

What LLMs Actually Are

An LLM is a prediction engine.

Not a reasoning engine. Not a database. Not a search system.

Given text, it predicts what comes next. That prediction runs over and over, token by token, until a full response is generated. The model learns these predictions from enormous amounts of human text: books, articles, code, research papers, websites.

Input:  "The capital of France is"
Model:  [predicts next token]
Output: "Paris"

Simple idea. The scale is what makes it work.

A token is roughly 3 to 4 characters of English text. "Hello, world!" is about 4 tokens. Everything the model processes and generates is counted in tokens. This affects cost, speed, and the limits we will cover shortly.

Why It Felt Like a Big Deal

For the first time, a machine could:

Hold a fluent conversation in any language
Write code that actually ran
Summarize a 50-page document in seconds
Explain complex topics to a non-expert
Answer questions across almost any domain

Where It Broke Down

Then people started building real products. The limitations became obvious fast.

Hallucination: The model predicts what is plausible, not what is true. It will state wrong facts with total confidence.

Knowledge cutoff: Training data has a date. Ask about last week and it guesses.

No memory: Every conversation starts blank. The model has no idea what you talked about yesterday.

No access to your data: Your company documents, your database, your internal systems. The model knows none of it.

No ability to act: It produces text. It cannot send an email, run a query, or update a record.

Ask a pure LLM: "What was Apple's stock price yesterday?" It will either refuse or make up a number.

It has no connection to live systems. It is a very smart autocomplete engine. Autocomplete alone does not run a business.

This limitation is what created the next stage.

Stage 2: RAG Changes the Game

The Core Idea

RAG stands for Retrieval-Augmented Generation. One sentence covers it:

Before generating a response, retrieve the relevant information and give it to the model.

Instead of relying only on training data, the system fetches fresh, relevant context at the moment of each query.

A Simple Way to Think About It

Pure LLM: A student answering an exam entirely from memory. Sometimes brilliant. Sometimes confidently wrong.

RAG: The same student, but allowed to open their notes before answering. Answers are grounded in actual sources.

The model did not get smarter. It got better information to work with.

How RAG Works

USER QUERY
    ↓
RETRIEVE relevant documents
(from a vector database, using semantic search)
    ↓
INJECT those documents into the prompt as context
    ↓
LLM generates an answer grounded in the retrieved content
    ↓
RESPONSE (accurate, with sources)

The Technology Behind Retrieval

Embeddings are what make semantic search work.

Documents are converted into vectors, which are lists of numbers that represent meaning mathematically. Similar meanings end up close together in vector space. "Car" and "automobile" are close. "Car" and "photosynthesis" are not.

When a user query arrives, it is also converted to a vector. The system finds the stored vectors nearest to that query vector. Those are the semantically relevant documents, retrieved and injected into context.

Common vector databases:

Database	Best For
Pinecone	Managed, production-ready
Weaviate	Open-source, rich query support
Chroma	Development and small-scale use
FAISS	Fast, local, no managed infrastructure

What RAG Unlocked

RAG became the foundation of a lot of serious AI products:

Enterprise knowledge assistants
Customer support bots grounded in actual policy
PDF and document Q&A
Internal search that surfaces the right document
Any system needing up-to-date or private data

What RAG Still Could Not Do

Retrieval solves the knowledge problem. It does not solve the action problem.

RAG can find the answer to "what is our refund policy?" It cannot process the refund. It can tell you flight options. It cannot book the ticket.

For that, a different capability was needed.

Stage 3: The Rise of AI Agents

The Core Shift

Traditional AI:

User asks → Model answers → Done

Agent:

User sets a goal → Agent plans → Agent uses tools →
Agent observes results → Agent decides next step →
Agent continues until goal is complete

Agents reason, plan, use tools, and execute multi-step workflows. They operate rather than just respond.

Tool Calling: How Agents Reach the Real World

An LLM by itself cannot search Google, call an API, write to a database, or run code. Tool calling extends the model's reach.

User: "Find the cheapest flights from Delhi to Singapore next month."

Agent Step 1: Call flight search API with parameters
Agent Step 2: Receive results
Agent Step 3: Sort and compare options
Agent Step 4: Summarize the three cheapest options

The model decides which tool to call, with what arguments, and what to do with the result. It manages the whole workflow.

What Agents Can Do

A capable AI agent can:

Browse websites and extract information
Write, execute, and debug code
Send emails and messages
Query and update databases
Call any API with proper credentials
Coordinate with other agents
Schedule and manage workflows

Frameworks That Made This Practical

Building agents from scratch is tedious. Frameworks handle the boilerplate:

LangChain / LangGraph - most widely used, graph-based orchestration
AutoGen - multi-agent conversations, good for collaborative tasks
CrewAI - role-based agent crews for structured workflows
OpenAI Agents SDK - native tool calling with built-in orchestration

Where Agents Break

More power introduced more failure modes:

Context overflow: Long agent runs fill the context window. Earlier instructions get lost. Accuracy drops.

Memory fragmentation: Without a coherent memory system, agents lose track of what they were doing.

Tool confusion: Too many tools and the model picks the wrong one or misuses it.

Hallucinated actions: The model invents results from tool calls it never actually made.

Runaway loops: No stop condition means the agent keeps going when it should have asked for clarification.

There was also a deeper infrastructure problem. Every agent integration was custom-built. Connecting to Slack required one connector. Google Drive required another. Salesforce required another. There was no standard. Scaling meant a growing stack of hand-built code.

That is what MCP was built to fix.

Stage 4: MCP, The Protocol That Standardizes Everything

The Problem Before MCP

Before November 2024, connecting an AI system to external tools meant:

Custom integration for every tool
Different formats for every API
No standard for how models discover what tools are available
No consistent way to pass context or results between systems

Every new data source required its own implementation. This was not an AI limitation. It was an infrastructure limitation.

What MCP Is

MCP stands for Model Context Protocol. It is a standard for connecting AI assistants to the systems where data lives, including content repositories, business tools, and development environments.

Anthropic announced MCP in November 2024 and open-sourced it on day one.

MCP defines a universal interface for:

Reading files and data sources
Executing functions and tools
Handling context and prompts
Coordinating between AI systems and external environments

The USB-C analogy is actually a good one here. Just as USB-C made it easy to connect any device to any peripheral, MCP makes it easier to connect any AI model to any data source or tool. One protocol, many integrations.

How It Works Architecturally

MCP servers expose three things:

Tools: actions the model can call
Resources: data the model can read
Prompts: templates for interaction

The model queries the server to discover what is available, then invokes tools in a structured, validated format.

Adoption

MCP did not take years to catch on. Since launch, it has been adopted by OpenAI, Microsoft, Google, and Cloudflare. The Python and JavaScript SDKs together see over 20 million weekly downloads. Over 13,000 MCP servers launched on GitHub in 2025 alone.

In December 2025, Anthropic donated MCP to the Agentic AI Foundation under the Linux Foundation, co-founded by Anthropic, Block, and OpenAI. It now sits alongside Kubernetes and PyTorch in that portfolio.

The Honest Limitations

MCP is not perfect. Security is a real concern.

Security researchers identified multiple issues with the protocol, including prompt injection, tool permissions that allow data exfiltration, and lookalike tools that can silently replace trusted ones.

The spec does not enforce audit trails, sandboxing, or verification. MCP solves the connectivity problem. Organizations deploying it at scale are responsible for building the security layer on top.

Context Engineering: The Layer Tying It All Together

Context engineering is what makes everything above work reliably in production.

Prompt engineering is writing a good instruction. Context engineering is designing the entire information environment the model operates in:

Memory: what the model remembers from previous interactions
Retrieval: what documents or data are fetched for each query
Tools: what actions are available and how they are described
History: how much of the conversation is included
System state: what the model knows about its current task
Workflow position: where in a multi-step process the model is

The most capable AI systems today are not just better models. They are better systems designed around the model.

Context engineering is what separates an agent that works in production from one that works in a demo.

What the Modern AI Stack Looks Like

A serious AI product in 2026 is a system, not just an API call:

User Interface
    ↓
Orchestration Layer (LangGraph, AutoGen, custom)
    ↓
Context Manager
├── Memory Layer (conversation history, user preferences)
├── Retrieval Layer (vector DB, semantic search)
└── State Manager (task progress, tool outputs)
    ↓
Tool Layer (via MCP or custom integrations)
├── Web search
├── Database queries
├── API calls
├── Code execution
└── File operations
    ↓
LLM (GPT-4o, Claude, Gemini, open-source)
    ↓
Response + Actions

Each layer solves a specific failure from the layer before it. Remove any layer and you reintroduce the problem it was built to solve.

Which Layer Do You Actually Need?

Do not over-engineer this.

A simple RAG pipeline is the right call for most document Q&A use cases. A complex agent adds coordination overhead you do not need if the task is just retrieval.

Add a layer only when the simpler system actually cannot meet the requirement.

What Comes Next

A few trends worth watching:

Long-term memory: Agents that remember your preferences across months, not just sessions.

Multi-agent collaboration: Networks of specialized agents coordinating on shared goals, where each handles one domain.

Deeper real-world execution: Tighter integration with operating systems and software, not just APIs.

Autonomous workflows: Agents that manage their own task queues without step-by-step human orchestration.

The bottleneck has moved. In 2020, it was model intelligence. In 2026, it is system design: how well you manage memory, retrieval, tool coordination, and context across a complex workflow.

The Real Takeaway

The biggest mistake people make is thinking the model is the entire product.

It is not.

Modern AI systems are architectures: memory, retrieval, orchestration, tool ecosystems, context managers, and execution environments wrapped around a model.

The future will not be decided only by which model is best.

It will be decided by which system is built best around it.

If you found this useful, I write about AI engineering weekly in my newsletter AI Engineering Simplified. No hype, just practical breakdowns.

[Boost]

Divy Yadav — Sun, 19 Apr 2026 17:51:48 +0000

Divy Yadav

Apr 14

7 Agent Harness Components Every AI Engineer Needs for Production

#agents #ai #architecture #softwareengineering

4 min read

[Boost]

Divy Yadav — Wed, 15 Apr 2026 08:45:26 +0000

Divy Yadav

Apr 14

7 Agent Harness Components Every AI Engineer Needs for Production

#agents #ai #architecture #softwareengineering

4 min read

7 Agent Harness Components Every AI Engineer Needs for Production

Divy Yadav — Tue, 14 Apr 2026 08:29:06 +0000

For a more in-depth guide

Medium: Link

The infrastructure layer your agents can't live without

Your agent just billed a user $38 on a single query.

Not because it did something complex.

Because it summarized the same document 47 times in a loop.

No crash. No alert. Just a growing invoice.

You check logs. The model worked exactly as expected.

The failure was everything around the model.

No memory. No state. No stop condition.

The system had no way to say:

"We’ve already done this."

That is the difference between a demo and production.

The gap nobody warns you about

Building a working agent is easy.

Call an LLM
Give it tools
Add a loop

20 lines of Python. Done.

Then you ship it.

Reality hits:

Tools return empty
Context overflows
Agents contradict themselves
Infinite retries start

Nothing breaks in demos.

Everything breaks in production.

The problem is not the model.

It is the harness.

The real definition

Agent = Model + Harness

Model → reasoning, decisions
Harness → execution, control, safety, memory

If you're not building the model, you're building the harness.

And that’s where most failures live.

The 7 components that actually matter

1. Control Loop

This is the heartbeat.

Without it → chatbot

With it → agent

while agent_is_running:
    response = call_model(context)

    if response.has_tool_calls:
        results = execute_tools(response.tool_calls)
        append_to_context(results)
        continue

    if response.is_final_answer:
        return response.content

    if step_count > MAX_STEPS:
        return "Task incomplete. Max steps reached."

Critical rule:

MAX_STEPS is non-negotiable.

No step limit = infinite billing.

2. State Management

Models are stateless.

You must track:

Session state

conversation history
tool outputs
current step

Persistent state

completed tasks
progress
processed files

Example:

{
  "task_id": "refactor-auth-module",
  "completed_files": ["auth.py"],
  "pending_files": ["routes.py"],
  "current_step": 3
}

Without this → agents repeat work endlessly.

3. Memory

State = what happened now

Memory = what survives later

Short-term

conversation history

Long-term

user preferences
past outcomes
domain knowledge

Typical flow:

Start:
  Load memory → inject into prompt

During:
  Maintain history

End:
  Summarize → store

Without memory → users feel your agent is dumb.

4. Tools (and the Bash Escape Hatch)

Tools convert language into action.

Bad tools:

vague descriptions
unclear usage

Good tools:

clear purpose
defined inputs/outputs
explicit usage rules

The real unlock: Bash

Instead of fixed tools → let agent create tools dynamically.

This is powerful.

Also dangerous.

So you need:

Sandboxing

isolated execution
no host access
safe parallel runs

Without sandbox → you are gambling.

5. Context Management

Silent killer.

Everything works... until it doesn’t.

Why?

Context fills up → important instructions get buried.

Solutions

1. Compaction

summarize old messages
keep system prompt intact

2. Truncation

limit tool outputs
store full data externally

3. Progressive loading

load tools only when needed

Rule:

Never lose the task definition.

6. Planning

Without planning → chaos.

Agent takes steps, but not the right ones.

Plan file pattern

task: Migrate database
steps:
  - Backup DB         [ ]
  - Run migration     [x]
  - Verify data       [ ]
current_step: 2

Each loop:

inject plan
update progress
verify step

Key concept: Self-verification

After each step:

run tests
validate output

Agents that verify → reliable

Agents that assume → break

Ralph Loop (important)

When context ends:

reload goal
continue from state

This enables long-running agents.

7. Error Handling

Reality is messy.

Things will fail.

You must define behavior for each failure.

Example logic

Tool fails:
  Retry (if temporary)
  Switch approach (if data issue)
  Escalate (if blocked)

Malformed output:
  Retry with constraints
  Fallback after 3 attempts

Loop detected:
  Interrupt execution

Low confidence:
  Send to human review

No error handling = hallucination or silent failure.

What actually happens in production

Example task:

"Summarize EU AI regulation news"

Flow:

Plan created
State initialized
Search executed
Articles fetched
Context managed
Summary generated
Verification step
Final output
Memory updated

The model writes.

The harness makes it reliable.

Common failure cases

Infinite loops

Fix → step limits + repetition detection

Tool misuse

Fix → better tool descriptions

Context overflow

Fix → compaction strategy

Hallucination with tools

Fix → enforce tool usage rules

Latency explosion

Fix → parallel tool execution

Hidden truth: Model ≠ independent

Modern agents are trained with harnesses.

Meaning:

Models adapt to specific tool patterns
Changing harness can reduce performance

Same model. Different harness. Different results.

When NOT to use agents

Be honest here.

Use deterministic pipelines when:

steps are fixed
outputs are predictable

Use humans when:

mistakes are costly

Avoid agents when:

task is structured and rule-based

If every step is predefined, an agent is overkill.

Where to start (practical order)

Control loop + step limit
State tracking
Small toolset
Error handling
Context management
Memory
Planning

Skip order → debug hell.

Final thought

The real test is not:

What happens when everything works?

It is:

What happens when things break?

Models will improve.

Harness design is what makes them usable.

The model is not your agent.

The harness is.

Google's TurboQuant: How They Cut LLM Memory by 6x Without Losing Accuracy

Divy Yadav — Fri, 27 Mar 2026 17:12:10 +0000

A plain-English breakdown of the Google Research paper that compresses KV cache by up to 6x with near-zero accuracy loss. No training. No calibration data. Just math.

Read the full indepth article on Medium: Link

Running large language models is not just expensive.

It is wasteful.

Every time you send a long prompt, the model stores massive amounts of intermediate data in something called the KV cache. This cache grows with every token. It quietly eats GPU memory, slows responses, and drives up inference costs.

Most compression solutions force a tradeoff. You either save memory or you keep accuracy. Pick one.

Google's TurboQuant breaks that tradeoff. It compresses the KV cache by up to 6x and, in several benchmarks, performs identically to the full-precision model.

That is a different kind of result. This post explains why, in plain English.

What Is the KV Cache?

Before anything else, you need to understand what TurboQuant is actually compressing.

When a language model processes text, it breaks input into tokens (roughly one token per word). For each token, it computes two sets of numbers: a Key and a Value. Together, these let the model decide how much attention to pay to each previous token when generating the next word.

Think of it this way. The Key is a label. The Value is the actual content. When the model wants to recall something, it scans all the Keys to find relevant ones, then reads the corresponding Values.

The KV cache stores all these Key-Value pairs for the entire conversation so the model does not recompute them on every new token.

The problem: these pairs are stored as 16-bit or 32-bit floating point numbers. A single layer of a large model can hold tens of millions of them. With hundreds of layers and thousands of tokens, the KV cache can consume gigabytes of GPU memory.

What quantization does: instead of storing each number at full precision, you store it at 3 or 4 bits. Smaller footprint. Faster lookups. But you lose accuracy in the rounding. The challenge is keeping that loss small.

What Vectors Actually Are (And Why This Matters)

TurboQuant is a vector quantization algorithm. If you have not worked with vectors in AI, here is the minimum you need.

A vector is an array of numbers. In AI, vectors encode meaning. The word "cat" might be represented as 1,536 numbers. Similar words (cat, kitten, feline) have similar vectors. Distant concepts (cat, telescope, inflation) have very different ones.

Vector quantization compresses those arrays by replacing precise values (like 0.7832194...) with approximate ones drawn from a small set of allowed values. The compression is lossy. The question is how much meaning you lose.

The naive approach groups a vector into small blocks, computes the centroid of each block, and stores which block each number belongs to.

The problem: every block needs its own scaling constant stored in full precision. Those constants add 1 to 2 extra bits per number, partially defeating the compression.

TurboQuant eliminates this overhead entirely. That is the core innovation.

The Old Problem: Overhead That Defeats the Point

Traditional vector quantization has a structural flaw that researchers have tolerated for decades.

When you quantize a block of numbers, you need to know the scale of that block to recover the original values. Numbers in the range 0 to 100 compress differently than numbers in the range 0 to 0.001. So you store the scale factor alongside the quantized data.

Those scale factors require full precision (16 or 32 bits) and must be stored per block. On a 4-bit scheme, adding even one extra bit of overhead per number is a 25% storage increase. It partially cancels out what you just compressed.

Most existing methods, including KIVI (a widely used KV cache quantizer), carry this overhead. It has been treated as unavoidable.

TurboQuant's two-stage design removes it by choosing a quantization strategy that does not need per-block scaling constants.

How TurboQuant Works: Two Stages

The clever part is not one idea. It is two ideas that work together.

Stage 1: PolarQuant

Standard quantization struggles because raw vectors have unpredictable distributions. Some dimensions have large values, others tiny ones. You need those scaling constants to manage this variation.

PolarQuant solves this by converting vectors into polar coordinates before quantizing.

Here is what that means. Instead of describing a point as "3 units east, 4 units north" (Cartesian), polar coordinates say "5 units total at a 37-degree angle." The distance from origin is the radius. The direction is the angle.

When you convert a high-dimensional vector into polar form, something useful happens. The radius captures signal strength. The angles capture direction and meaning. And the angles follow a predictable, concentrated distribution that fits neatly into a fixed grid.

Because the grid boundaries are predictable from geometry alone, you do not need separate scaling constants per block. The structure of the data handles it.

But there is a catch.

PolarQuant introduces a subtle bias in inner product estimation. Inner products are what the model uses when computing attention scores: how much should token 47 attend to token 12? A biased estimate means those scores drift slightly from what a full-precision model would compute. At low bit-widths (1 to 2 bits), this drift becomes significant.

Stage 2 fixes it.

Stage 2: QJL (Quantized Johnson-Lindenstrauss)

After Stage 1 quantizes most of the information, a small residual error remains: the gap between the quantized vector and the original.

TurboQuant takes that residual and applies a 1-bit correction using the Quantized Johnson-Lindenstrauss (QJL) transform:

Multiply the residual vector by a random matrix
Take only the sign of each resulting number (+1 or -1)
Store one bit per dimension

That sounds lossy. And it is. But mathematically, this specific random projection preserves inner product relationships in expectation. The QJL correction is unbiased: on average, it accurately corrects for the bias Stage 1 introduces, without adding any overhead for scaling constants.

Put together:

Stage 1 handles most of the compression and eliminates per-block overhead
Stage 2 uses 1 bit per dimension to correct the bias Stage 1 introduces
The result is near-optimal by information theory standards

The paper proves mathematically that TurboQuant's distortion is within a factor of approximately 2.7 of the theoretical minimum for any quantization algorithm. You cannot do significantly better without violating information theory.

The Results (This Is Where It Gets Interesting)

TurboQuant is not just theoretically sound. It actually holds up in benchmarks.

KV Cache Compression on LongBench

On long-context tasks (question answering, summarization, code generation):

At 3.5 bits: TurboQuant matches full precision exactly. That is already 4.5x smaller.
At 2.5 bits (~6x compression): the performance drop is minimal (50.06 to 49.44 on LongBench, a 1.2% drop).

Compared to KIVI and PolarQuant, TurboQuant either matches or outperforms at every bit level.

Needle-in-a-Haystack

This test checks if a model can locate one specific piece of information buried inside a massive document.

Method	Score
TurboQuant	0.997
PolarQuant	0.995
KIVI	0.981
PyramidKV	0.895
SnapKV	0.858
Full precision	~1.000

TurboQuant works up to 100K tokens with near-perfect recall.

Attention Computation Speed

Up to 8x faster attention computation compared to standard 32-bit models on H100 GPUs. Not just smaller. Also faster.

Vector Search

This one surprised me.

Method	Indexing Time
Product Quantization	Seconds to minutes
RabitQ	Even slower
TurboQuant	~0.001 seconds

And despite that indexing speed, TurboQuant still achieves better recall accuracy than both methods.

No training. No calibration data. Near-instant indexing.

What "Data-Oblivious" Actually Means

Most quantization methods are data-dependent. They analyze a sample of your data, learn the distribution, and build a custom codebook optimized for that distribution.

This is why methods like Product Quantization and GPTQ take seconds or minutes to build an index. It also means they degrade silently when your data distribution shifts.

TurboQuant is data-oblivious. It applies the same random rotation to any input, then uses pre-computed quantization centroids based on the mathematical properties of rotated vectors. No learning. No calibration data. No dataset-specific tuning.

This matters for three real scenarios:

KV cache quantization: tokens arrive in real time. There is no opportunity to pre-analyze the data.

Live vector search systems: new vectors are added continuously. Re-indexing periodically is expensive. Data drift degrades data-dependent quantizers silently.

Production systems with shifting distributions: TurboQuant sidesteps the re-index-or-degrade tradeoff entirely.

Where This Actually Applies

Long-context LLM inference

Running Llama-3.1 or Gemini on long documents (contracts, research papers, codebases) is expensive because the KV cache grows linearly with context length. TurboQuant compresses it 4.5x to 6x without accuracy loss. The same hardware now handles roughly 4 to 6x longer contexts at the same cost.

RAG and vector databases

If you run a retrieval-augmented generation system with Qdrant, Pinecone, Weaviate, or pgvector, those databases store embedding vectors. Compressing them with TurboQuant means smaller index storage, faster similarity lookups, near-zero indexing time for new documents, and no need to retrain the quantizer when your data changes.

Edge and on-device deployment

Phones, laptops, embedded systems have strict memory budgets. KV cache compression is sometimes the difference between a model that fits in device memory and one that does not. TurboQuant gives 4x to 6x compression with no accuracy tradeoff.

Semantic search at scale

Google's own writeup explicitly mentions this as a primary use case. Building and querying vector indices over billions of embeddings requires minimal memory and fast lookups. TurboQuant's near-zero indexing time and recall numbers make it a strong option here.

What TurboQuant Does Not Do

Being honest about scope matters.

It is not for weight quantization. TurboQuant is designed for activations (KV cache values) and vector search indices. It is not a replacement for GPTQ or AWQ, which target model weights.

It has been tested on specific models. The KV cache experiments used Llama-3.1-8B-Instruct and Ministral-7B-Instruct. Results may vary on significantly different architectures, though the algorithm is theoretically model-agnostic.

There is no plug-and-play PyPI package yet. The paper is on arXiv (arXiv:2504.19874) and was presented at ICLR 2026. Google Research has an implementation, but as of publication there is no standalone production library. The algorithms are described clearly enough that a research team could implement from the paper.

The 2.5-bit case does trade some accuracy for compression. The 1.2% drop on LongBench may or may not be acceptable depending on your use case. That is a decision that depends on your application, not something TurboQuant resolves for you.

The Mental Model

Traditional quantization stores numbers in small blocks and needs extra overhead bits to record the scale of each block. Those overhead bits eat into the compression benefit.

TurboQuant avoids the problem by rotating vectors into a geometry where the scale is already known from the structure of the data, not from per-block constants. A 1-bit residual correction from QJL handles the bias this rotation introduces.

The result is near-optimal compression with provable guarantees, no training required, and near-zero preprocessing time.

What it does	Result
KV cache compression	4.5x to 6x with no accuracy loss
Needle-in-haystack	0.997 (matches full precision)
LongBench general tasks	Matches full precision at 3.5 bits
Attention computation	Up to 8x faster vs 32-bit on H100
Vector search indexing	~0.001 seconds vs 37 to 600 seconds
Training required	None

Final Thought

For a long time, memory in large language models has been treated as a necessary cost. If you want long context, you pay for it.

TurboQuant challenges that assumption.

It shows that with the right mathematical approach, you can significantly reduce memory usage without sacrificing performance. Not by building a smarter model. By compressing the existing one more honestly.

The next wave of AI progress may not come from bigger models. It might come from making existing systems faster, cheaper, and more efficient.

If that is true, TurboQuant is not just an optimization. It is a direction worth paying attention to.

Paper: TurboQuant: Online Vector Quantization with Near-optimal Distortion — arXiv:2504.19874

If you found this useful, drop a reaction or share it with someone working on LLM inference or vector search. And if you have tried KV cache quantization in production,
I would genuinely like to hear what you ran into.

I Tried Vector Search on Molecules. Here Is What Actually Happened.

Divy Yadav — Fri, 27 Mar 2026 09:07:17 +0000

In this article, I'll walk you through how I built a robust molecular similarity search system using ChemBERTa, RDKit, Qdrant, and what I actually learned along the way.

Canonical URL: https://medium.com/towards-artificial-intelligence/i-tried-vector-search-on-molecules-heres-what-happened-7391b755efe4

TL;DR
I wanted to see if vector search could work on molecules the same way it works on text. It can. I used ChemBERTa to embed SMILES strings into 768-dim vectors, indexed them in Qdrant with molecular property metadata, and ran similarity search with payload filters applied during retrieval. The system surfaced structurally similar candidates that fingerprint-based search missed. This post walks through every step, including where it breaks.

GitHub: github.com/dvy246/qdrant
Stack: Python 3.10+, RDKit, ChemBERTa, Qdrant, FastAPI, Streamlit

Why I Built This

I had been spending a lot of time with vector databases and embedding-based search. Every example I came across was about text: document search, FAQ retrieval, chatbot memory.

At some point, I asked myself a question that seemed a bit strange at the time: Can I apply the same thing to molecules?

Molecules can be written as text strings. SMILES (Simplified Molecular Input Line Entry System) is a notation that encodes molecular structure as a string. CC(=O)Oc1ccccc1C(=O)O is aspirin. If you can write a molecule as a string, you can feed it to a transformer. And if you can get a dense vector out of a transformer, you can search it with a vector database.

I had been reading about ChemBERTa, a transformer trained specifically on SMILES strings. Qdrant had just shipped a clean Python client with native payload filtering. I thought: what happens if I connect these two things?

So I built a small experiment and ran it.

Here is what the system does:

Downloads and caches the ZINC-250k dataset (2000 drug-like molecules by default)
Validates and canonicalizes each SMILES string using RDKit, computing a heuristic toxicity proxy per molecule
Embeds each molecule into a 768-dimensional vector using ChemBERTa
Indexes those vectors in Qdrant with molecular property metadata as searchable payload
Retrieves similar molecules using cosine similarity with optional filters on MW, LogP, and toxicity
Serves results through a FastAPI endpoint or a Streamlit UI

The Problem with Fingerprints

Most molecule search pipelines use fingerprint similarity. The basic idea: take a molecule, generate a Morgan fingerprint (a fixed-length bit vector that records which molecular fragments are present), compute a Tanimoto score against a query molecule (overlap between the two bit vectors, 0 to 1), and rank by that score.

For obvious structural matches, this works well. Fast, simple, interpretable.

But when I was testing it on a small set of compounds, I kept seeing it miss molecules that looked chemically related when I inspected them manually. Not enough fragment overlap for the score to reflect that relationship. The system just did not surface them.

*That gap is what pushed me toward trying something different.
*
The core issue is that compressing a molecule into a fixed-length bit vector loses information. Different substructures can hash to the same bit position. That collision is gone before you run a query. You are searching for a lossy representation of the molecule, not the molecule itself.

This gets particularly visible in what chemists call scaffold hopping: two molecules that act on the same biological target through structurally different cores. Their fingerprints can look completely different even when they are functionally related. Tanimoto has no way to catch that.

There is also the activity cliff problem. Adding a single methyl group (one carbon with three hydrogens) can completely flip how a molecule behaves biologically. Tanimoto treats that as a single-bit difference with no chemical context.

And at scale, Tanimoto is a full linear scan. With millions of compounds and thousands of queries, that is billions of pairwise comparisons. It requires dedicated engineering to keep running at a reasonable speed.

None of this was a deal-breaker for a small experiment. But it was enough to make me curious whether an approach that learned something about molecular structure, rather than just counting fragments, could surface what fingerprints miss.

The Idea: Embeddings Over Fragments

Instead of hashing molecules into fragment bit vectors, I wanted to try converting each molecule into a dense vector that captured its structural and physicochemical character in continuous space. The bet: similar molecules would land close together in that space even when their fingerprints look different.

ChemBERTa handles that. It is a RoBERTa-based transformer trained on SMILES strings from the ZINC database. During training, it learns to predict masked tokens from surrounding context, the same pretraining approach BERT uses for text. Because SMILES encodes molecular structure in a consistent grammar, the model picks up structural patterns across millions of molecules. The result is a 768-number vector per molecule.

One thing to be clear about: these vectors are structural estimates, not measures of biological activity. The base model was trained without any assay data. You are measuring learned structural closeness, not functional equivalence. That limitation matters and comes up again in the limitations section.

Once molecules are represented as vectors, approximate nearest neighbor search replaces the linear scan. That is what Qdrant handles.

Pipeline Overview

Five stages, each doing one thing:

SMILES strings
    │
    ▼
RDKit validation and canonicalization
    │
    ▼
ChemBERTa embedding (Hugging Face Transformers)
    │
    ▼
Qdrant vector indexing (upsert with metadata payloads)
    │
    ▼
Similarity search API (FastAPI) or Web UI (Streamlit)

Raw SMILES come in first. RDKit validates and standardizes them. Anything that does not parse gets dropped before it reaches the model, and more comes in invalid than you might expect from a curated source.

Validated SMILES go to ChemBERTa (seyonec/ChemBERTa-zinc-base-v1). Mean pooling on the final hidden state gives a 768-dimensional vector per molecule. That vector, along with metadata like molecular weight and LogP, gets stored as a point in Qdrant. HNSW handles the indexing. At query time, the same model embeds the query, Qdrant searches by cosine similarity, and results can be filtered by molecular properties during retrieval, not after.

Why Qdrant

I had one core requirement: apply numeric filters on molecular properties during retrieval, not after it. That ruled out more options than I expected.

Pinecone was the easiest to get running. Basic vector search worked fine. The problem: combining similarity search with numeric range filters on molecular weight and LogP produced inconsistent behavior.

Milvus gave good raw search performance. The problem: local installation required spinning up etcd and MinIO just to run a dev experiment. Not acceptable for something I was iterating on quickly.

Weaviate had a clean schema system I liked. The problem: chaining multiple numeric range filters with a vector search query felt awkward. The semantics were not intuitive enough for this use case.

Qdrant solved the specific problems the others did not.

The key difference is how it handles filters. In this kind of search you rarely query by vector alone. You want to say: find molecules similar to this one, but only where molecular weight is under 500 and LogP is below 5. Qdrant applies those constraints natively during HNSW graph traversal, not after. Post-filtering sounds like the same thing but it is not. If your filter is restrictive, you silently end up with fewer than top-k results because candidates are removed after the search finishes. That was not acceptable for how I needed this to work.

The in-memory client (QdrantClient(":memory:")) was a big deal for development. I could build and test the entire pipeline on a laptop without Docker or any external service. Moving to a persistent Qdrant instance later was one line. That made the iteration loop fast.

When to Use Embeddings vs. Fingerprints

If you are working with a well-characterized scaffold series where structure-activity relationships are already mapped, ECFP4 and Tanimoto are probably still the right tool. Fast, interpretable, no model required. Adding an embedding pipeline for that use case adds complexity without giving much back.

Embeddings start to make sense when you are trying to surface candidates that do not share obvious substructure with your query, or when your library is large enough that a linear Tanimoto scan becomes a bottleneck.

ChemBERTa is not the only option. Mol2Vec produces lighter embeddings and runs faster on CPU. GNN-based models can encode 3D molecular geometry, which matters if shape drives binding for your target. I went with ChemBERTa because it takes SMILES as direct input and is well-supported on Hugging Face. For this kind of experiment, that simplicity was worth more than the extra capability of a more complex model.

Similarity Metric

Before indexing, every vector gets L2 normalized. Once vectors are unit length, cosine similarity and dot product give identical results.

I went with cosine similarity for one practical reason: the scores are easy to read. Cosine scores sit between -1 and 1, so a score of 0.95 versus 0.72 tells you something concrete when you are staring at retrieval results. When you are debugging why certain molecules surface and others do not, readable scores matter.

Environment Setup

Python 3.10+ required. All dependencies from PyPI.

# Core cheminformatics
pip install rdkit

# Transformer framework, tokenizer, and PyTorch
pip install transformers torch

# Vector database client (requires Qdrant server v1.10+ for the Query API)
pip install "qdrant-client>=1.10.0"

# API or UI (pick one or both)
pip install fastapi uvicorn
pip install streamlit

Note: The PyPI package is rdkit, not rdkit-pypi (older name). The codebase uses X | Y union type syntax which requires Python 3.10+. ChemBERTa runs on CPU for smaller datasets but GPU inference is significantly faster above 100k molecules.

Verify before starting:

import rdkit
from rdkit import Chem
from transformers import AutoTokenizer, AutoModel
from qdrant_client import QdrantClient

print(f"RDKit version: {rdkit.__version__}")
print("All imports successful.")

Building It: Step by Step

Step 1: Load the Dataset and Validate SMILES

SMILES strings from real databases are messy. The same molecule can have multiple valid representations. Some entries contain salts or stereoisomer annotations. Others are just invalid. Feeding any of that directly to a transformer produces garbage embeddings or silent failures.

I went with ZINC-250k, a widely used benchmark subset of drug-like molecules from the ZINC database. It is the same source ChemBERTa was pre-trained on, which makes it a natural fit.

# data_loader.py
from __future__ import annotations

import csv
import logging
from pathlib import Path

from rdkit import Chem
from molsearch.config import DATA_CACHE_DIR, DATASET_SIZE

logger = logging.getLogger(__name__)

ZINC_250K_URL = "https://raw.githubusercontent.com/aspuru-guzik-group/chemical_vae/master/models/zinc_properties/250k_rndm_zinc_drugs_clean_3.csv"
ZINC_FILENAME = "zinc_250k.csv"

_FALLBACK_SMILES: list[str] = [
    "CC(=O)Oc1ccccc1C(=O)O",          # Aspirin
    "CC(C)Cc1ccc(cc1)C(C)C(=O)O",     # Ibuprofen
    "CC(=O)Nc1ccc(O)cc1",             # Acetaminophen
    "O=C(O)Cc1ccccc1Nc1c(Cl)cccc1Cl", # Diclofenac
    "COc1ccc2cc(CC(C)C(=O)O)ccc2c1",  # Naproxen
]


def load_dataset(
    max_molecules: int = DATASET_SIZE,
) -> tuple[list[str], list[float | None]]:
    """
    Load molecules from the ZINC-250k dataset.
    Downloads on first call, caches locally for subsequent runs.
    Toxicity scores are None — computed dynamically in molecule_processor.py.
    Falls back to built-in list if download fails.
    """
    cache_dir = Path(DATA_CACHE_DIR)
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_dir / ZINC_FILENAME

    if not cache_path.exists():
        try:
            import urllib.request
            logger.info("Downloading ZINC-250k from %s ...", ZINC_250K_URL)
            urllib.request.urlretrieve(ZINC_250K_URL, str(cache_path))
        except Exception:
            logger.warning("Download failed; using %d fallback molecules", len(_FALLBACK_SMILES))
            smiles = _FALLBACK_SMILES[:max_molecules]
            return smiles, [None] * len(smiles)

    seen: set[str] = set()
    valid_smiles: list[str] = []

    with open(cache_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            raw = row.get("smiles", "").strip()
            if not raw:
                continue
            mol = Chem.MolFromSmiles(raw)
            if mol is None or mol.GetNumAtoms() == 0:
                continue
            canonical = Chem.MolToSmiles(mol)
            if canonical in seen:
                continue
            seen.add(canonical)
            valid_smiles.append(canonical)
            if len(valid_smiles) >= max_molecules:
                break

    toxicity_scores: list[float | None] = [None] * len(valid_smiles)
    logger.info("Dataset ready: %d molecules", len(valid_smiles))
    return valid_smiles, toxicity_scores

DATASET_SIZE defaults to 2000, controlled via MOLSEARCH_DATASET_SIZE env var. Toxicity scores are all None here because ZINC has no toxicity labels. They get computed per molecule in the validation step.

Validation and canonicalization runs over every SMILES. RDKit checks that the string parses to a valid molecule and converts it to canonical form so there is exactly one text representation per structure regardless of how it was originally written.

Heads up on stereochemistry: RDKit canonicalization does not preserve stereocenters if they are not explicitly defined. Two enantiomers can end up with identical canonical SMILES and identical vectors. If chirality matters for your use case, handle it before this step.

from __future__ import annotations

import logging
import math

from rdkit import Chem
from rdkit.Chem import Descriptors, rdMolDescriptors

logger = logging.getLogger(__name__)


def compute_toxicity_proxy(mol: Chem.Mol) -> float:
    """
    Heuristic toxicity estimate from RDKit descriptors.

    NOT a real toxicity prediction. Rule-based proxy using properties
    commonly associated with toxicity risk. Returns float 0.0–1.0.
    """
    mw     = Descriptors.MolWt(mol)
    logp   = Descriptors.MolLogP(mol)
    hbd    = Descriptors.NumHDonors(mol)
    hba    = Descriptors.NumHAcceptors(mol)
    tpsa   = Descriptors.TPSA(mol)
    n_arom = rdMolDescriptors.CalcNumAromaticRings(mol)

    mw_score   = max(0.0, min((mw - 350)  / 450, 1.0))
    logp_score = max(0.0, min((logp - 3)  / 5,   1.0))
    hbd_score  = max(0.0, min((hbd - 2)   / 5,   1.0))
    hba_score  = max(0.0, min((hba - 5)   / 10,  1.0))
    arom_score = min(n_arom / 6, 1.0)
    tpsa_score = max(0.0, min((75 - tpsa) / 75,  1.0))

    raw = (
        0.25 * mw_score
        + 0.25 * logp_score
        + 0.15 * hbd_score
        + 0.15 * hba_score
        + 0.10 * arom_score
        + 0.10 * tpsa_score
    )
    return round(max(0.0, min(raw, 1.0)), 3)


def validate_and_canonicalize(
    smiles: str,
    toxicity_score: float | None = None,
) -> dict | None:
    """
    Validate a SMILES string and return canonical form with descriptors.
    Returns None if the SMILES is invalid.
    """
    normalized = smiles.strip()
    if not normalized:
        return None

    mol = Chem.MolFromSmiles(normalized)
    if mol is None:
        return None

    num_atoms = mol.GetNumAtoms()
    if num_atoms == 0 or mol.GetNumBonds() == 0:
        return None

    canonical_smiles = Chem.MolToSmiles(mol)

    # Round-trip check: re-parse canonical SMILES and verify atom count
    verify_mol = Chem.MolFromSmiles(canonical_smiles)
    if verify_mol is None or verify_mol.GetNumAtoms() != num_atoms:
        return None

    payload = {
        "smiles":            canonical_smiles,
        "molecular_weight":  round(Descriptors.MolWt(mol), 2),
        "logp":              round(Descriptors.MolLogP(mol), 2),
        "num_h_donors":      Descriptors.NumHDonors(mol),
        "num_h_acceptors":   Descriptors.NumHAcceptors(mol),
        "tpsa":              round(Descriptors.TPSA(mol), 2),
    }

    if toxicity_score is not None:
        if not math.isfinite(toxicity_score):
            raise ValueError("toxicity_score must be a finite float")
        payload["toxicity_score"] = float(toxicity_score)
    else:
        payload["toxicity_score"] = compute_toxicity_proxy(mol)

    return payload

Two things worth calling out. When toxicity_score is None (always the case for ZINC data), the function falls through to compute_toxicity_proxy. Every molecule in the index ends up with a toxicity_score in its payload, so filtered searches on that field always have something to work with. The proxy is a weighted combination of six RDKit descriptors. It is a deterministic heuristic, not a toxicity model.

The round-trip check after canonicalization catches edge cases where RDKit generates a canonical form it then cannot parse back. Rare, but it happens.

Step 2: Generate Molecular Embeddings with ChemBERTa

ChemBERTa is a RoBERTa model pre-trained on SMILES strings from ZINC. Checkpoint: seyonec/ChemBERTa-zinc-base-v1 on Hugging Face.

The model does not come with a dedicated pooling layer. To get a single vector per molecule, you pool the token-level outputs yourself. I use mean pooling — averaging the output vectors across all non-padding tokens. This is the standard approach for RoBERTa-based models because the [CLS] token (a special first token added to every input in BERT-style models) was not trained with a sentence-level objective. Taking just that token gives unstable representations. Averaging across all tokens works better.

from __future__ import annotations

import logging
import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer

logger = logging.getLogger(__name__)

MODEL_NAME = "seyonec/ChemBERTa-zinc-base-v1"
VECTOR_DIM  = 768
BATCH_SIZE  = 32


def _get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


class MoleculeEmbedder:
    """Generates L2-normalized ChemBERTa embeddings from SMILES."""

    def __init__(self, model_name: str = MODEL_NAME):
        self.device    = _get_device()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model     = AutoModel.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()
        self.vector_dim = self.model.config.hidden_size
        if self.vector_dim != VECTOR_DIM:
            raise ValueError(
                f"Model hidden size ({self.vector_dim}) does not match "
                f"VECTOR_DIM ({VECTOR_DIM})"
            )
        logger.info("Loaded %s on %s", model_name, self.device)

    def _mean_pool(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        mask_expanded = (
            attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
        )
        sum_hidden = torch.sum(hidden_states * mask_expanded, dim=1)
        sum_mask   = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
        return sum_hidden / sum_mask

    def embed(
        self,
        smiles_list: list[str],
        batch_size: int = BATCH_SIZE,
    ) -> np.ndarray:
        """
        Embed SMILES into dense vectors.
        Returns (n, vector_dim) float32 array, L2-normalized.
        Oversized SMILES (>400 chars) are skipped per-molecule and filled with np.nan.
        """
        n = len(smiles_list)
        if n == 0:
            return np.empty((0, self.vector_dim), dtype=np.float32)

        result = np.empty((n, self.vector_dim), dtype=np.float32)

        for start in range(0, n, batch_size):
            end   = min(start + batch_size, n)
            batch = smiles_list[start:end]

            # Handle oversized SMILES per-molecule, not per-batch
            oversized = [i for i, s in enumerate(batch) if len(s) > 400]
            if oversized:
                for idx in oversized:
                    logger.warning(
                        "SMILES at index %d too long (%d chars), skipping",
                        start + idx, len(batch[idx]),
                    )
                    result[start + idx] = np.nan
                safe_indices = [i for i in range(len(batch)) if i not in oversized]
                if not safe_indices:
                    continue
                batch = [batch[i] for i in safe_indices]
            else:
                safe_indices = list(range(len(batch)))

            encoded = self.tokenizer(
                batch, padding=True, truncation=False, return_tensors="pt"
            )
            if encoded["input_ids"].shape[1] > 512:
                logger.error("Batch %d–%d: context limit exceeded", start, end)
                result[start:end] = np.nan
                continue

            encoded = {k: v.to(self.device) for k, v in encoded.items()}
            with torch.no_grad():
                outputs = self.model(**encoded)

            embeddings = self._mean_pool(
                outputs.last_hidden_state, encoded["attention_mask"]
            )
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            batch_result = embeddings.cpu().numpy()

            for out_i, safe_i in enumerate(safe_indices):
                result[start + safe_i] = batch_result[out_i]

        return result

A few implementation details worth knowing.

truncation=False is deliberate. Silently truncating a SMILES string produces an embedding that only represents part of the molecule — that is worse than no embedding at all. Instead, any SMILES over 400 characters is flagged and skipped per-molecule. If the tokenized batch still exceeds 512 tokens after that, the whole batch gets rejected. Failed slots are filled with np.nan so downstream code can detect the problem rather than silently work with bad data.

Result array is pre-allocated with np.empty rather than appending batch results to a list. For large datasets, np.vstack temporarily doubles memory because it copies everything into a new contiguous block. Pre-allocation avoids that.

CPU performance: expect somewhere around 100–300 molecules per second depending on SMILES length. Above 100k molecules, GPU inference saves a meaningful amount of time.

Step 3: Index Embeddings in Qdrant

With valid molecules and their embeddings ready, the next step is storing them in Qdrant. Each point holds a vector and a JSON payload with canonical SMILES and all the descriptors from Step 1.

from __future__ import annotations

import uuid
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

COLLECTION_NAME    = "molecules"
VECTOR_DIM         = 768
UPSERT_BATCH_SIZE  = 1000
USE_PERSISTENT_QDRANT = False
QDRANT_HOST        = "localhost"
QDRANT_PORT        = 6333


def get_qdrant_client() -> QdrantClient:
    """In-memory for development, persistent for production."""
    if USE_PERSISTENT_QDRANT:
        client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT, timeout=10)
        client.get_collections()  # verify connection
        return client
    return QdrantClient(":memory:")


def create_collection(client: QdrantClient) -> None:
    """Non-destructive: reuse existing collection if vector size matches."""
    if client.collection_exists(collection_name=COLLECTION_NAME):
        info = client.get_collection(collection_name=COLLECTION_NAME)
        size = getattr(info.config.params.vectors, "size", None)
        if size is not None and size != VECTOR_DIM:
            raise ValueError(
                f"Existing collection has vector size {size}, expected {VECTOR_DIM}"
            )
        return
    client.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=VECTOR_DIM, distance=Distance.COSINE),
    )


def _smiles_to_uuid(smiles: str) -> str:
    """Deterministic UUID from canonical SMILES. Safe for incremental upserts."""
    namespace = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
    return str(uuid.uuid5(namespace, smiles))


def upsert_molecules(
    client: QdrantClient,
    molecules: list[dict],
    embeddings: np.ndarray,
    batch_size: int = UPSERT_BATCH_SIZE,
) -> None:
    """Upsert molecule vectors and payloads in batches. Skips nan/inf/zero vectors."""
    n = len(molecules)
    valid_mask = (
        ~np.isnan(embeddings).any(axis=1)
        & ~np.isinf(embeddings).any(axis=1)
    )
    valid_mask &= np.linalg.norm(embeddings, axis=1) > 0.0

    for start in range(0, n, batch_size):
        end = min(start + batch_size, n)
        points = [
            PointStruct(
                id=_smiles_to_uuid(mol["smiles"]),
                vector=emb.tolist(),
                payload=mol,
            )
            for i, (mol, emb) in enumerate(
                zip(molecules[start:end], embeddings[start:end]), start=start
            )
            if valid_mask[i]
        ]
        if points:
            client.upsert(collection_name=COLLECTION_NAME, points=points)


# Create payload indexes for filtered search
from qdrant_client.models import PayloadSchemaType

for field_name in ["molecular_weight", "logp", "toxicity_score"]:
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name=field_name,
        field_schema=PayloadSchemaType.FLOAT,
    )

Point IDs are generated using uuid.uuid5 derived from canonical SMILES. Re-indexing the same molecule just overwrites the existing entry. Re-running the indexing script without wiping the collection is safe.

The payload index step matters at scale. Without it, filtered searches slow down noticeably once the collection grows, because Qdrant has to scan payload fields without an index.

Step 4: Similarity Search with Hybrid Ranking

Given a query SMILES, embed it with the same ChemBERTa model and search Qdrant for nearest neighbors.

from __future__ import annotations

from rdkit.Chem import rdFingerprintGenerator, DataStructs
from qdrant_client.models import Filter, FieldCondition, Range


def search_similar_molecules(
    query_smiles: str,
    embedder: MoleculeEmbedder,
    client: QdrantClient,
    collection_name: str = COLLECTION_NAME,
    top_k: int = 5,
    mw_max: float | None = None,
    logp_max: float | None = None,
    toxicity_max: float | None = None,
) -> list[dict]:
    try:
        embeddings = embedder.embed([query_smiles])
        if np.any(np.isnan(embeddings)):
            raise ValueError("query vector generation failed (nan)")
        query_vector = embeddings[0].tolist()
    except Exception as exc:
        logger.error("Failed to embed query SMILES: %s", exc)
        raise RuntimeError(f"Query embedding failed: {exc}") from exc

    conditions = []
    if mw_max is not None:
        conditions.append(FieldCondition(key="molecular_weight", range=Range(lte=mw_max)))
    if logp_max is not None:
        conditions.append(FieldCondition(key="logp", range=Range(lte=logp_max)))
    if toxicity_max is not None:
        conditions.append(FieldCondition(key="toxicity_score", range=Range(lte=toxicity_max)))

    query_filter = Filter(must=conditions) if conditions else None
    fetch_limit  = top_k * 5  # fetch wider pool before reranking

    results = None
    for attempt in range(3):
        try:
            results = client.query_points(
                collection_name=collection_name,
                query=query_vector,
                query_filter=query_filter,
                limit=fetch_limit,
                with_payload=True,
                timeout=30,
            )
            break
        except Exception as exc:
            if attempt == 2:
                logger.error("Qdrant query failed after 3 attempts: %s", exc)
                return []

    if results is None:
        return []

    query_mol = Chem.MolFromSmiles(query_smiles)
    if query_mol is not None:
        generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
        query_fp  = generator.GetFingerprint(query_mol)
    else:
        query_fp = None

    hits = []
    for point in results.points:
        payload    = point.payload or {}
        hit_smiles = payload.get("smiles", "")
        tanimoto   = 0.0

        if query_fp is not None and hit_smiles:
            hit_mol = Chem.MolFromSmiles(hit_smiles)
            if hit_mol is not None:
                hit_fp   = generator.GetFingerprint(hit_mol)
                tanimoto = DataStructs.TanimotoSimilarity(query_fp, hit_fp)

        # TODO: make fusion weights configurable
        fused_score = 0.5 * point.score + 0.5 * tanimoto

        hits.append({
            "smiles":           hit_smiles,
            "score":            round(point.score, 4),
            "tanimoto_score":   round(tanimoto, 4),
            "fused_score":      round(fused_score, 4),
            "molecular_weight": payload.get("molecular_weight", 0.0),
            "logp":             payload.get("logp", 0.0),
            "toxicity_score":   (
                payload.get("toxicity_score")
                if isinstance(payload.get("toxicity_score"), (int, float)) else None
            ),
        })

    hits.sort(key=lambda x: x["fused_score"], reverse=True)
    return hits[:top_k]

The search retrieves top_k * 5 candidates from Qdrant before reranking. HNSW is approximate, so fetching a wider pool gives a better shot at returning the actual best matches after fusion.

The final ranking combines two signals: ChemBERTa cosine similarity (learned structural patterns) and Tanimoto fingerprint score (explicit substructure overlap). The equal 50/50 weight split is a starting point. The TODO in the code is intentional — for any real use case, tune those weights against your own data.

Step 5: FastAPI Service

# api_server.py
from __future__ import annotations

import asyncio
import logging
import math
from contextlib import asynccontextmanager
from functools import partial

from fastapi import FastAPI, HTTPException, Response
from pydantic import BaseModel, Field
from rdkit import Chem

from molsearch.config import MAX_SMILES_LENGTH
from molsearch.data_loader import load_dataset
from molsearch.embedder import MoleculeEmbedder
from molsearch.molecule_processor import process_smiles_batch
from molsearch.qdrant_indexer import (
    check_system_health, collection_exists_and_populated,
    create_collection, create_payload_indexes,
    get_qdrant_client, search_similar_molecules, upsert_molecules,
)

logger = logging.getLogger(__name__)

_embedder: MoleculeEmbedder | None = None
_client = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _embedder, _client
    try:
        _embedder = MoleculeEmbedder()
    except Exception as exc:
        logger.error("Failed to load embedder: %s", exc)
        _embedder = None

    try:
        _client = get_qdrant_client()
        create_collection(_client)
        create_payload_indexes(_client)
        if not collection_exists_and_populated(_client):
            if _embedder is not None:
                smiles_list, toxicity_scores = load_dataset()
                molecules   = process_smiles_batch(smiles_list, toxicity_scores=toxicity_scores)
                embeddings  = _embedder.embed([m["smiles"] for m in molecules])
                upsert_molecules(_client, molecules, embeddings)
                logger.info("Indexed %d molecules.", len(molecules))
        else:
            logger.info("Collection already populated — skipping indexing.")
    except Exception as exc:
        logger.error("Failed to initialize Qdrant: %s", exc)
        _client = None

    yield
    _client   = None
    _embedder = None


app = FastAPI(title="Molecule Similarity Search API", lifespan=lifespan)


class SearchRequest(BaseModel):
    smiles:       str         = Field(..., min_length=1, max_length=MAX_SMILES_LENGTH)
    top_k:        int         = Field(default=5, ge=1, le=100)
    mw_max:       float | None = Field(default=None, ge=0, allow_inf_nan=False)
    logp_max:     float | None = Field(default=None, allow_inf_nan=False)
    toxicity_max: float | None = Field(default=None, ge=0, allow_inf_nan=False)


class MoleculeHit(BaseModel):
    smiles:           str
    score:            float
    molecular_weight: float
    logp:             float
    toxicity_score:   float | None = None
    tanimoto_score:   float | None = None
    fused_score:      float | None = None


class SearchResponse(BaseModel):
    query_smiles:    str
    canonical_smiles: str
    results:         list[MoleculeHit]


@app.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
    if _embedder is None or _client is None:
        raise HTTPException(status_code=503, detail="Service not initialized")

    query_smiles = request.smiles.strip()
    if not query_smiles:
        raise HTTPException(status_code=400, detail="empty smiles")

    mol = Chem.MolFromSmiles(query_smiles)
    if mol is None or mol.GetNumAtoms() == 0:
        raise HTTPException(status_code=400, detail="invalid smiles")

    canonical = Chem.MolToSmiles(mol)
    if len(canonical) > 400:
        raise HTTPException(status_code=400, detail={
            "error":   "too_long",
            "message": "smiles exceeds length limit",
            "length":  len(canonical),
        })

    loop = asyncio.get_running_loop()
    try:
        hits = await loop.run_in_executor(
            None,
            partial(
                search_similar_molecules,
                query_smiles=canonical,
                embedder=_embedder,
                client=_client,
                top_k=request.top_k,
                mw_max=request.mw_max,
                logp_max=request.logp_max,
                toxicity_max=request.toxicity_max,
            ),
        )
    except RuntimeError as exc:
        raise HTTPException(status_code=500, detail=f"Search failed: {exc}") from exc

    return SearchResponse(
        query_smiles=request.smiles,
        canonical_smiles=canonical,
        results=[MoleculeHit(**h) for h in hits],
    )


@app.get("/health")
def health(response: Response):
    health_status = check_system_health(_embedder, _client)
    response.status_code = 200 if health_status["status"] == "ok" else 503
    return health_status

ChemBERTa inference is CPU-bound. Running it directly inside the FastAPI endpoint would block the async event loop. run_in_executor hands the transformer inference and Qdrant search to a separate thread, keeping the event loop free for other requests.

uvicorn molsearch.api_server:app --host 0.0.0.0 --port 8000 --reload

Test it:

curl -X POST http://localhost:8000/search \
  -H "Content-Type: application/json" \
  -d '{"smiles": "CC(=O)Oc1ccccc1C(=O)O", "top_k": 3}'

Step 6: Streamlit UI

# streamlit_app.py
from __future__ import annotations

import streamlit as st
from rdkit import Chem
from rdkit.Chem import Descriptors, Draw

from molsearch.data_loader import load_dataset
from molsearch.embedder import MoleculeEmbedder
from molsearch.molecule_processor import process_smiles_batch
from molsearch.qdrant_indexer import (
    collection_exists_and_populated, create_collection,
    create_payload_indexes, get_qdrant_client,
    search_similar_molecules, upsert_molecules,
)

st.set_page_config(page_title="Molecule Similarity Search", layout="wide")
st.title("Molecule Similarity Search")
st.caption("Powered by ChemBERTa embeddings and Qdrant vector search")


@st.cache_resource
def load_resources():
    embedder = MoleculeEmbedder()
    client   = get_qdrant_client()
    create_collection(client)
    create_payload_indexes(client)
    if not collection_exists_and_populated(client):
        smiles_list, toxicity_scores = load_dataset()
        molecules  = process_smiles_batch(smiles_list, toxicity_scores=toxicity_scores)
        embeddings = embedder.embed([m["smiles"] for m in molecules])
        upsert_molecules(client, molecules, embeddings)
    return embedder, client


embedder, client = load_resources()

st.sidebar.header("Search Parameters")
query_smiles = st.sidebar.text_input(
    "SMILES string",
    value="CC(=O)Oc1ccccc1C(=O)O",
    help="Enter a valid SMILES string",
)
top_k = st.sidebar.slider("Number of results", min_value=1, max_value=20, value=5)

use_mw_filter = st.sidebar.checkbox("Filter by molecular weight")
mw_filter     = (
    st.sidebar.number_input("Max molecular weight", value=500.0, step=50.0, min_value=0.0)
    if use_mw_filter else None
)

use_logp_filter = st.sidebar.checkbox("Filter by LogP")
logp_filter     = (
    st.sidebar.number_input("Max LogP", value=5.0, step=0.5)
    if use_logp_filter else None
)

use_tox_filter  = st.sidebar.checkbox("Filter by toxicity")
toxicity_filter = (
    st.sidebar.number_input("Max toxicity score", value=0.5, step=0.1, min_value=0.0)
    if use_tox_filter else None
)

search_clicked = st.sidebar.button("Search", type="primary", use_container_width=True)

if search_clicked:
    mol = Chem.MolFromSmiles(query_smiles.strip())
    if mol is None or mol.GetNumAtoms() == 0:
        st.error(f"Invalid SMILES: {query_smiles}")
    else:
        canonical_query = Chem.MolToSmiles(mol)
        col_query, col_info = st.columns([1, 2])
        with col_query:
            st.subheader("Query Molecule")
            st.image(Draw.MolToImage(mol, size=(300, 300)), caption=canonical_query)
        with col_info:
            st.subheader("Query Info")
            st.metric("Molecular Weight", f"{Descriptors.MolWt(mol):.2f}")
            st.metric("LogP", f"{Descriptors.MolLogP(mol):.2f}")
        st.divider()

        results = search_similar_molecules(
            query_smiles=canonical_query, embedder=embedder, client=client,
            top_k=top_k, mw_max=mw_filter, logp_max=logp_filter, toxicity_max=toxicity_filter,
        )
        st.subheader(f"Top {len(results)} Similar Molecules")

        if not results:
            st.info("No results found. Try relaxing the filters.")
        else:
            for i, hit in enumerate(results):
                with st.container():
                    c1, c2 = st.columns([1, 2])
                    with c1:
                        hit_mol = Chem.MolFromSmiles(hit["smiles"])
                        if hit_mol:
                            st.image(Draw.MolToImage(hit_mol, size=(250, 250)))
                    with c2:
                        st.markdown(
                            f"**Rank {i + 1}** | "
                            f"Fused: **{hit['fused_score']}** "
                            f"(Tanimoto: {hit['tanimoto_score']}, Latent: {hit['score']})"
                        )
                        st.code(hit["smiles"], language=None)
                        toxicity_text = (
                            f"{hit['toxicity_score']:.3f}"
                            if isinstance(hit.get("toxicity_score"), (int, float)) else "n/a"
                        )
                        st.write(f"MW: {hit['molecular_weight']} | LogP: {hit['logp']} | Toxicity: {toxicity_text}")
                    st.divider()
else:
    st.info("Enter a SMILES string in the sidebar and click Search.")

streamlit run src/molsearch/streamlit_app.py

@st.cache_resource keeps the 400MB ChemBERTa model from reloading on every interaction. Without it, adjusting any slider would reload the entire model. The load_dataset() call inside load_resources() also only runs once per server session.

End-to-End Standalone Script

If you want to run the full pipeline without the modular project structure:

"""
molecule_search_pipeline.py
End-to-end molecular similarity search.
Usage: python molecule_search_pipeline.py
"""
from __future__ import annotations

import uuid
import numpy as np
import torch
from rdkit import Chem
from rdkit.Chem import Descriptors
from transformers import AutoTokenizer, AutoModel
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

MODEL_NAME      = "seyonec/ChemBERTa-zinc-base-v1"
COLLECTION_NAME = "molecules"
BATCH_SIZE      = 32


def process_smiles(smiles_list: list[str]) -> list[dict]:
    results = []
    for smi in smiles_list:
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            continue
        results.append({
            "smiles":           Chem.MolToSmiles(mol),
            "molecular_weight": round(Descriptors.MolWt(mol), 2),
            "logp":             round(Descriptors.MolLogP(mol), 2),
            "num_h_donors":     Descriptors.NumHDonors(mol),
            "num_h_acceptors":  Descriptors.NumHAcceptors(mol),
            "tpsa":             round(Descriptors.TPSA(mol), 2),
        })
    return results


def load_model(model_name: str = MODEL_NAME):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model     = AutoModel.from_pretrained(model_name)
    model.eval()
    return tokenizer, model, model.config.hidden_size


def embed_smiles(smiles_list, tokenizer, model, vector_dim=768) -> np.ndarray:
    if not smiles_list:
        return np.empty((0, vector_dim), dtype=np.float32)
    all_embeddings = []
    for i in range(0, len(smiles_list), BATCH_SIZE):
        batch   = smiles_list[i:i + BATCH_SIZE]
        encoded = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**encoded)
        mask     = encoded["attention_mask"].unsqueeze(-1).expand(outputs.last_hidden_state.size()).float()
        summed   = torch.sum(outputs.last_hidden_state * mask, dim=1)
        counted  = torch.clamp(mask.sum(dim=1), min=1e-9)
        embeddings = summed / counted
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
        all_embeddings.append(embeddings.cpu().numpy())
    return np.vstack(all_embeddings)


def index_molecules(client, molecules, embeddings, vector_dim):
    if client.collection_exists(collection_name=COLLECTION_NAME):
        client.delete_collection(collection_name=COLLECTION_NAME)
    client.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=vector_dim, distance=Distance.COSINE),
    )
    points = [
        PointStruct(
            id=str(uuid.uuid5(uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8"), mol["smiles"])),
            vector=emb.tolist(),
            payload=mol,
        )
        for mol, emb in zip(molecules, embeddings)
    ]
    client.upsert(collection_name=COLLECTION_NAME, points=points)
    print(f"  Indexed {len(points)} molecules.")


def search(client, query_smiles, tokenizer, model, vector_dim=768, top_k=5):
    from qdrant_client.models import Filter, FieldCondition, Range
    mol = Chem.MolFromSmiles(query_smiles)
    if mol is None:
        raise ValueError(f"Invalid query SMILES: {query_smiles}")
    canonical  = Chem.MolToSmiles(mol)
    query_vec  = embed_smiles([canonical], tokenizer, model, vector_dim)[0].tolist()
    results    = client.query_points(
        collection_name=COLLECTION_NAME,
        query=query_vec,
        limit=top_k,
        with_payload=True,
    )
    return results.points


def main():
    from molsearch.data_loader import load_dataset

    print("Loading ZINC-250k dataset...")
    raw_smiles, _ = load_dataset()
    print(f"  {len(raw_smiles)} molecules loaded.\n")

    print("Step 1: Processing SMILES...")
    molecules = process_smiles(raw_smiles)
    print(f"  {len(molecules)} valid molecules.\n")

    print("Step 2: Loading model and generating embeddings...")
    tokenizer, model, vector_dim = load_model()
    embeddings = embed_smiles([m["smiles"] for m in molecules], tokenizer, model, vector_dim)
    print(f"  Embedding matrix shape: {embeddings.shape}\n")

    print("Step 3: Indexing in Qdrant...")
    client = QdrantClient(":memory:")
    index_molecules(client, molecules, embeddings, vector_dim)

    print("\nStep 4: Searching...")
    query = "CC(=O)Oc1ccccc1C(=O)O"  # Aspirin
    print(f"  Query: Aspirin ({query})")
    hits = search(client, query, tokenizer, model, vector_dim, top_k=5)
    print(f"\n  {'Rank':<6}{'SMILES':<50}{'Score':<10}")
    print(f"  {'-' * 66}")
    for i, hit in enumerate(hits, 1):
        print(f"  {i:<6}{hit.payload['smiles']:<50}{hit.score:<10.4f}")


if __name__ == "__main__":
    main()

Expected output:

Step 1: Processing SMILES...
  2000 valid molecules.

Step 2: Loading model and generating embeddings...
  Embedding matrix shape: (2000, 768)

Step 3: Indexing in Qdrant...
  Indexed 2000 molecules.

Step 4: Searching...
  Query: Aspirin (CC(=O)Oc1ccccc1C(=O)O)

  Rank  SMILES                                            Score
  ------------------------------------------------------------------
  1     CC(=O)Oc1ccccc1C(=O)O                             1.0000
  2     OC(=O)c1ccccc1O                                   0.95xx
  3     OC(=O)c1ccccc1                                    0.93xx
  ...

Scaling Considerations

The example above indexes 2000 molecules. Once the pipeline was working end to end, the obvious question was: what changes at millions?

HNSW Tuning

from qdrant_client.models import HnswConfigDiff, OptimizersConfigDiff

client.create_collection(
    collection_name="molecules_prod",
    vectors_config=VectorParams(size=768, distance=Distance.COSINE),
    hnsw_config=HnswConfigDiff(
        m=16,             # edges per node (default 16)
        ef_construct=128, # search depth during construction (default 100)
    ),
    optimizers_config=OptimizersConfigDiff(
        indexing_threshold=20000,  # triggers around ~6.7k vectors at 768-dim float32
    ),
)

Parameter	Low (fast, lower recall)	High (slower, higher recall)
`m`	8	32
`ef_construct`	64	256
Query `ef`	64	256

For similarity search, recall matters more than shaving milliseconds. A missed candidate is a missed result.

Set ef at query time:

from qdrant_client import models

results = client.query_points(
    collection_name=COLLECTION_NAME,
    query=query_vec,
    search_params=models.SearchParams(hnsw_ef=128),
    limit=top_k,
    with_payload=True,
)

Embedding at Scale

For datasets over 100k molecules:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

encoded = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors="pt")
encoded = {k: v.to(device) for k, v in encoded.items()}

Other optimizations that help: increase batch size to 128 or 256, use torch.amp.autocast("cuda") for mixed-precision inference, pre-compute embeddings and store in Parquet files so you do not rerun the model on every experiment iteration.

Persistent Deployment

docker run -p 6333:6333 -v $(pwd)/qdrant_storage:/qdrant/storage qdrant/qdrant

client = QdrantClient(host="localhost", port=6333)

Dev to production is one line. The rest of the code does not change.

Honest Limitations

No experiment writeup is complete without the part where you explain where it breaks.

The toxicity score is not a toxicity model. compute_toxicity_proxy produces a number between 0 and 1 based on six RDKit descriptors. It is useful for demonstrating filtered retrieval. Do not make safety decisions with it.

ChemBERTa operates in 1D and 2D. It reads SMILES strings and learns topological patterns. It has no concept of 3D geometry. If shape complementarity matters for your target, this pipeline alone is not enough. You would need to layer in 3D-aware models or shape-based screening.

Similarity scores are structural estimates. The base model was trained on ZINC using masked language modeling with no bioactivity labels. Two molecules that look close in vector space can behave completely differently in an assay. Vector retrieval is a starting point, not a replacement for experimental validation.

Implementation-level gotchas:

Canonical SMILES can differ across RDKit versions. Pin your version and keep it consistent between indexing and query environments.
HNSW is approximate. If you need high recall, increase ef at query time.
Salt-containing SMILES like [Na+].[Cl-] will be parsed as a whole fragment including the counterion. Add rdkit.Chem.SaltRemover.SaltRemover to preprocessing if that matters for your data.
The token limit for ChemBERTa is 512. The production implementation rejects SMILES over 400 characters before tokenization. The standalone script uses truncation=True, which silently cuts long inputs. Know which behavior you are running.

Where to Take This Next

Full code: github.com/dvy246/qdrant

Three directions worth thinking about.

Fine-tuning. The base model has no activity labels, so the similarity it measures is purely structural. Fine-tuning ChemBERTa on specific assay data with a contrastive objective — so active molecules pull together in vector space — is the most direct way to make retrieval useful for a specific target. Without that, you are doing structural exploration, not activity prediction.

Automated ingestion. Connect Qdrant updates into your compound registration workflow so new molecules get indexed as they come in rather than in periodic batch jobs. Qdrant's deterministic UUID approach makes this safe. Re-indexing the same molecule just overwrites.

Target-aware search. Add protein pocket embeddings alongside molecular embeddings so retrieval accounts for both the molecule and the target. That starts to look more like a real virtual screening tool.

Conclusion

The experiment worked, and honestly it worked better than I expected.

The pipeline loads 2000 molecules from ZINC-250k, validates and canonicalizes each SMILES with RDKit (computing a heuristic toxicity score per molecule while doing it), converts each structure into a 768-dimensional vector using ChemBERTa, stores that in Qdrant with molecular property metadata, and retrieves structurally similar candidates using cosine similarity. Filters on molecular weight, LogP, and toxicity score apply during retrieval rather than after it. Final ranking combines embedding similarity with a Tanimoto fingerprint score so you get both the learned structural space and explicit substructure overlap in one result.

This does not replace fingerprint search. For exact scaffold matching, Tanimoto is faster, simpler, and more interpretable. What this gives you is a way to find candidates that share structural and physicochemical patterns even when their fingerprints look different. That is the case where fingerprint search fails.

The hard parts were not the transformer or the vector database. They were the data cleaning, the SMILES canonicalization edge cases, and deciding when to trust the similarity scores. Those are the things worth understanding before scaling this up.

If you are a developer exploring vector search beyond text, molecules turn out to be a surprisingly good playground. The data is structured, the search requirements are concrete, and the gap between what fingerprints can do and what embeddings can do is real enough to make the experiment worthwhile.