That was my advice to a colleague recently when they asked about a classification problem. Who fine-tunes BERT anymore? Haven’t decoder models eaten the entire NLP landscape?

The look I got back was… skeptical. And it stuck with me.

I’ve been deep in LLM-land for a few years now. When your daily driver can architect systems, write production code, and reason through problems better than most junior devs, you start reaching for it reflexively. Maybe my traditional ML instincts had atrophied.

So I decided to actually test my assumptions instead of just vibing on them.

I ran 32 experiments pitting small instruction-tuned LLMs against good old BERT and DeBERTa. I figured I’d just be confirming what I already believed, that these new decoder models would obviously crush the ancient encoders.

I was wrong.

The results across Gemma 2B, Qwen 0.5B/1.5B, BERT-base, and DeBERTa-v3 were… not what I expected. If you’re trying to decide between these approaches for classification, you might want to actually measure things instead of assuming the newer model is better.

All the code is on GitHub if you want to run your own experiments.

Experiment Setup

What I Tested

BERT Family (Fine-tuned)

• BERT-base-uncased (110M parameters)
• DeBERTa-v3-base (184M parameters)

Small LLMs

• Qwen2-0.5B-Instruct
• Qwen2.5-1.5B-Instruct
• Gemma-2-2B-it

For the LLMs, I tried two approaches:

1. Zero-shot - Just prompt engineering, no training
2. Few-shot (k=5) - Include 5 examples in the prompt

Tasks

Four classification benchmarks ranging from easy sentiment to adversarial NLI:

Methodology

For anyone who wants to reproduce this or understand what “fine-tuned” and “zero-shot” actually mean here:

BERT/DeBERTa Fine-tuning:

• Standard HuggingFace Trainer with AdamW optimizer
• Learning rate: 2e-5, batch size: 32, epochs: 3
• Max sequence length: 128 tokens
• Evaluation on validation split (GLUE test sets don’t have public labels)

LLM Zero-shot:

• Greedy decoding (temperature=0.0) for deterministic outputs
• Task-specific prompts asking for single-word classification labels
• No examples in context—just instructions and the input text

LLM Few-shot (k=5):

• Same as zero-shot, but with 5 labeled examples prepended to each prompt
• Examples randomly sampled from training set (stratified by class)

All experiments used a fixed random seed (99) for reproducibility. Evaluation metrics are accuracy on the validation split. Hardware: RunPod instance with RTX A4500 (20GB VRAM), 20GB RAM, 5 vCPU.

I’d forgotten how pretty text-only land can be. When you spend most of your time in IDEs and notebooks, SSH-ing into a headless GPU box and watching nvitop do its thing feels almost meditative.

Results

Let’s dive into what actually happened:

DeBERTa-v3 wins most tasks—but not all

DeBERTa hit 94.8% on SST-2, 80.9% on RTE, and 82.6% on BoolQ. For standard classification with decent training data, the fine-tuned encoders still dominate.

On ANLI—the hardest benchmark, specifically designed to fool models—Gemma few-shot actually beats DeBERTa (47.8% vs 47.4%). It’s a narrow win, but it’s a win on the task that matters most for robustness.

Zero-shot LLMs actually beat BERT-base

The LLMs aren’t losing to BERT—they’re losing to DeBERTa. Qwen2.5-1.5B zero-shot hit 93.8% on SST-2, beating BERT-base’s 91.5%. Same story on RTE (78.7% vs 61.0%) and BoolQ (Gemma’s 80.9% vs BERT’s 71.5%). For models running purely on prompts with zero training? I’m calling it a win.

Few-shot is a mixed bag

Adding examples to the prompt doesn’t always help.

On RTE, Qwen2.5-1.5B went from 78.7% zero-shot down to 53.4% with few-shot. On SST-2, it dropped from 93.8% to 89.0%. But on ANLI, few-shot helped significantly—Gemma jumped from 36.1% to 47.8%, enough to beat DeBERTa.

Few-shot helps on harder tasks where examples demonstrate the thought process, but can confuse models on simpler pattern matching tasks where they already “get it.” Sometimes examples add noise instead of signal.

BERT Goes Brrrr

Okay, so the accuracy gap isn’t huge. Maybe I could still justify using an LLM?

Then I looked at throughput:

BERT is ~20x faster.

BERT processes 277 samples per second. Gemma-2-2B manages 12. If you’re classifying a million documents, that’s one hour vs a full day.

Encoders process the whole sequence in one forward pass. Decoders generate tokens autoregressively, even just to output “positive” or “negative”.

Note on LLM latency: These numbers use max_length=256 for tokenization. When I bumped it to max_length=2048, latency jumped 8x—from 57ms to 445ms per sample for Qwen-0.5B. Context window scales roughly linearly with inference time. For short classification tasks, keep it short or make it dynamic.

Try It Yourself

These models struggled on nuanced reviews. Can you do better? Try classifying some of the trickiest examples from my experiments:

When LLMs Make Sense

Despite the efficiency gap, there are cases where small LLMs are the right choice:

Zero Training Data

If you have no labeled data, LLMs win by default. Zero-shot Qwen2.5-1.5B at 93.8% on SST-2 is production-ready without a single training example. You can’t fine-tune BERT with zero examples.

Rapidly Changing Categories

If your categories change frequently (new product types, emerging topics), re-prompting an LLM takes seconds. Re-training BERT requires new labeled data, training time, validation, deployment. The iteration cycle matters.

Explanations with Predictions

LLMs can provide reasoning: “This review is negative because the customer mentions ‘defective product’ and ‘waste of money.’” BERT gives you a probability. Sometimes you need the story, not just the number.

Low Volume

If you’re processing 100 support tickets a day, throughput doesn’t matter. The 20x speed difference is irrelevant when you’re not hitting any resource constraints.

When BERT Still Wins

High-Volume Production Systems

If you’re classifying millions of items daily, BERT’s 20x throughput advantage matters. That’s a job finishing in an hour vs. running all day.

Well-Defined, Stable Tasks

Sentiment analysis. Spam detection. Topic classification. If your task definition hasn’t changed since 2019, fine-tuned BERT is proven and stable. No need to fix what isn’t broken.

You Have Training Data

With a few thousand labeled examples, fine-tuned DeBERTa will beat small LLMs. It’s a dedicated specialist vs. a generalist. Specialization still works.

Latency Matters

Real-time classification in a user-facing app where every millisecond counts? BERT’s parallel processing wins. LLMs can’t compete on speed.

Limitations

Before you @ me on Twitter—yes, I know this isn’t the final word. Some caveats:

I only tested small LLMs. Kept everything under 2B parameters to fit comfortably on a 20GB GPU. Bigger models like Llama-3-8B or Qwen-7B would probably do better, but then the efficiency comparison becomes even more lopsided. You’re not beating BERT’s throughput with a 7B model.

Generic prompts. I used straightforward prompts without heavy optimization. Task-specific prompt engineering could boost LLM performance. DSPy-style optimization would probably help too—but that’s another blog post.

Four benchmarks isn’t everything. There are plenty of classification scenarios I didn’t test. Your domain might be different. Measure, don’t assume.

Conclusion

So, can small LLMs beat BERT at classification?

Sometimes, and on the hardest task, they actually do. Gemma few-shot edges out DeBERTa on adversarial NLI, the benchmark specifically designed to break models.

DeBERTa-v3 still wins 3 out of 4 tasks when you have training data. And BERT’s efficiency advantage is real—~20x faster throughput matters when you’re processing millions of documents and paying for compute.

Zero-shot LLMs aren’t just a parlor trick either. Qwen2.5-1.5B hits 93.8% on sentiment with zero training examples—that’s production-ready without a single label. For cold-start problems, rapidly changing domains, or when you need explanations alongside predictions, they genuinely work.

Hopefully this gives some actual data points for making that call instead of just following the hype cycle.

All the code is on GitHub. Go run your own experiments.

Related reading

• A Production Framework for LLM Feature Evaluation — where classification fits among the LLM features that actually ship.
• RAG: From Context Injection to Knowledge Integration — the other pattern everyone reaches for, and its architectural limits.

Surely I’ve made some embarrassing mistakes here. Don’t just tell me—tell everyone! Share this post on your favorite social media with your corrections :)

Everyone Loves pgvector (in theory)

If you’ve spent any time in the vector search space over the past year, you’ve probably read blog posts explaining why pgvector is the obvious choice for your vector database needs. The argument goes something like this: you already have Postgres, vector embeddings are just another data type, why add complexity with a dedicated vector database when you can keep everything in one place?

It’s a compelling story. And like most of the AI influencer bullshit that fills my timeline, it glosses over the inconvenient details.

I’m not here to tell you pgvector is bad. It’s not. It’s a useful extension that brings vector similarity search to Postgres. But after spending some time trying to build a production system on top of it, I’ve learned that the gap between “works in a demo” and “scales in production” is… significant.

Nobody’s actually run this in production

What bothers me most: the majority of content about pgvector reads like it was written by someone who spun up a local Postgres instance, inserted 10,000 vectors, ran a few queries, and called it a day. The posts are optimistic, the benchmarks are clean, and the conclusions are confident.

They’re also missing about 80% of what you actually need to know.

I’ve read through   dozens of these posts.

What they don’t tell you is what happens when you actually try to run this in production.

Picking an index (there are no good options)

Let’s start with indexes, because this is where the tradeoffs start.

pgvector gives you two index types: IVFFlat and HNSW. The blog posts will tell you that HNSW is newer and generally better, which is… technically true but deeply unhelpful.

IVFFlat

IVFFlat (Inverted File with Flat quantization) partitions your vector space into clusters. During search, it identifies the nearest clusters and only searches within those.

The good:

• Lower memory footprint during index creation
• Reasonable query performance for many use cases
• Index creation is faster than HNSW

The bad:

• Requires you to specify the number of lists (clusters) upfront
• That number significantly impacts both recall and query performance
• The commonly recommended formula (rows / 1000) is a starting point at best
• Recall can be… disappointing depending on your data distribution
• New vectors get assigned to existing clusters, but clusters don’t rebalance without a full rebuild

Image source: IVFFlat or HNSW index for similarity search? by Simeon Emanuilov

HNSW

HNSW (Hierarchical Navigable Small World) builds a multi-layer graph structure for search.

The good:

• Better recall than IVFFlat for most datasets
• More consistent query performance
• Scales well to larger datasets

The bad:

• Significantly higher memory requirements during index builds
• Index creation is slow—painfully slow for large datasets
• The memory requirements aren’t theoretical; they are real, and they’ll take down your database if you’re not careful

Image source: IVFFlat or HNSW index for similarity search? by Simeon Emanuilov

None of the blogs mention that building an HNSW index on a few million vectors can consume 10+ GB of RAM or more (depending on your vector dimensions and dataset size). On your production database. While it’s running. For potentially hours.

Real-time search is basically impossible

In a typical application, you want newly uploaded data to be searchable immediately. User uploads a document, you generate embeddings, insert them into your database, and they should be available in search results. Simple, right?

How index updates actually work

When you insert new vectors into a table with an index, one of two things happens:

1. IVFFlat: The new vectors are inserted into the appropriate clusters based on the existing structure. This works, but it means your cluster distribution gets increasingly suboptimal over time. The solution is to rebuild the index periodically. Which means downtime, or maintaining a separate index and doing an atomic swap, or accepting degraded search quality.

2. HNSW: New vectors are added to the graph structure. This is better than IVFFlat, but it’s not free. Each insertion requires updating the graph, which means memory allocation, graph traversals, and potential lock contention.

Neither of these is a deal-breaker in isolation. But here’s what happens in practice: you’re inserting vectors continuously throughout the day. Each insertion is individually cheap, but the aggregate load adds up. Your database is now handling your normal transactional workload, analytical queries, AND maintaining graph structures in memory for vector search.

Handling new inserts

Let’s say you’re building a document search system. Users upload PDFs, you extract text, generate embeddings, and insert them. The user expects to immediately search for that document.

Here’s what actually happens:

With no index: The insert is fast, the document is immediately available, but your searches do a full sequential scan. This works fine for a few thousand documents. At a few hundred thousand? Your searches start taking seconds. Millions? Good luck.

With IVFFlat: The insert is still relatively fast. The vector gets assigned to a cluster. But whoops, a problem. Those initial cluster assignments were based on the data distribution when you built the index. As you add more data, especially if it’s not uniformly distributed, some clusters get overloaded. Your search quality degrades. You rebuild the index periodically to fix this, but during the rebuild (which can take hours for large datasets), what do you do with new inserts? Queue them? Write to a separate unindexed table and merge later?

With HNSW: The graph gets updated on each insert through incremental insertion, which sounds great. But updating an HNSW graph isn’t free—you’re traversing the graph to find the right place to insert the new node and updating connections. Each insert acquires locks on the graph structure. Under heavy write load, this becomes a bottleneck. And if your write rate is high enough, you start seeing lock contention that slows down both writes and reads.

The operational reality

Here’s the real nightmare: you’re not just storing vectors. You have metadata—document titles, timestamps, user IDs, categories, etc. That metadata lives in other tables (or other columns in the same table). You need that metadata and the vectors to stay in sync.

In a normal Postgres table, this is easy—transactions handle it. But when you’re dealing with index builds that take hours, keeping everything consistent gets complicated. For IVFFlat, periodic rebuilds are basically required to maintain search quality. For HNSW, you might need to rebuild if you want to tune parameters or if performance has degraded.

The problem is that index builds are memory-intensive operations, and Postgres doesn’t have a great way to throttle them. You’re essentially asking your production database to allocate multiple (possibly dozens) gigabytes of RAM for an operation that might take hours, while continuing to serve queries.

You end up with strategies like:

• Write to a staging table, build the index offline, then swap it in (but now you have a window where searches miss new data)
• Maintain two indexes and write to both (double the memory, double the update cost)
• Build indexes on replicas and promote them
• Accept eventual consistency (users upload documents that aren’t searchable for N minutes)
• Provision significantly more RAM than your “working set” would suggest

None of these are “wrong” exactly. But they’re all workarounds for the fact that pgvector wasn’t really designed for high-velocity real-time ingestion.

Pre- vs. Post-Filtering (or: why you need to become a query planner expert)

Okay but let’s say you solve your index and insert problems. Now you have a document search system with millions of vectors. Documents have metadata—maybe they’re marked as draft, published, or archived. A user searches for something, and you only want to return published documents.

1SELECT * FROM documents
2WHERE status = 'published'
3ORDER BY embedding <-> query_vector
4LIMIT 10;

Simple enough. But now you have a problem: should Postgres filter on status first (pre-filter) or do the vector search first and then filter (post-filter)?

This seems like an implementation detail. It’s not. It’s the difference between queries that take 50ms and queries that take 5 seconds. It’s also the difference between returning the most relevant results and… not.

Pre-filter works great when the filter is highly selective (1,000 docs out of 10M). It works terribly when the filter isn’t selective—you’re still searching millions of vectors.

Post-filter works when your filter is permissive. Here’s where it breaks: imagine you ask for 10 results with LIMIT 10. pgvector finds the 10 nearest neighbors, then applies your filter. Only 3 of those 10 are published. You get 3 results back, even though there might be hundreds of relevant published documents slightly further away in the embedding space.

The user searched, got 3 mediocre results, and has no idea they’re missing way better matches that didn’t make it into the initial k=10 search.

You can work around this by fetching more vectors (say, LIMIT 100) and then filtering, but now:

• You’re doing way more distance calculations than needed
• You still don’t know if 100 is enough
• Your query performance suffers
• You’re guessing at the right oversampling factor

With pre-filter, you avoid this problem, but you get the performance problems I mentioned. Pick your poison.

Multiple filters

Now add another dimension: you’re filtering by user_id AND category AND date_range.

1SELECT * FROM documents
2WHERE user_id = 'user123'
3  AND category = 'technical'
4  AND created_at > '2024-01-01'
5ORDER BY embedding <-> query_vector
6LIMIT 10;

What’s the right strategy now?

• Apply all filters first, then search? (Pre-filter)
• Search first, then apply all filters? (Post-filter)
• Apply some filters first, search, then apply remaining filters? (Hybrid)
• Which filters should you apply in which order?

The planner will look at table statistics, index selectivity, and estimated row counts and come up with a plan. That plan will probably be wrong, or at least suboptimal, because the planner’s cost model wasn’t built for vector similarity search.

And it gets worse: you’re inserting new vectors throughout the day. Your index statistics are outdated. The plans get increasingly suboptimal until you ANALYZE the table. But ANALYZE on a large table with millions of rows takes time and resources. And it doesn’t really understand vector data distribution in a meaningful way—it can tell you how many rows match user_id = 'user123', but not how clustered those vectors are in the embedding space, which is what actually matters for search performance.

Workarounds

You end up with hacks: query rewriting for different user types, partitioning your data into separate tables, CTE optimization fences to force the planner’s hand, or just fetching way more results than needed and filtering in application code.

None of these are sustainable at scale.

What vector databases do

Dedicated vector databases have solved this. They understand the cost model of filtered vector search and make intelligent decisions:

• Adaptive strategies: Some databases dynamically choose pre-filter or post-filter based on estimated selectivity
• Configurable modes: Others let you specify the strategy explicitly when you know your data distribution
• Specialized indexes: Some build indexes that support efficient filtered search (like filtered HNSW)
• Query optimization: They track statistics specific to vector operations and optimize accordingly

OpenSearch’s k-NN plugin, for example, lets you specify pre-filter or post-filter behavior. Pinecone automatically handles filter selectivity. Weaviate has optimizations for common filter patterns.

With pgvector, you get to build all of this yourself. Or live with suboptimal queries. Or hire a Postgres expert to spend weeks tuning your query patterns.

Hybrid search? Build it yourself

Oh, and if you want hybrid search—combining vector similarity with traditional full-text search—you get to build that yourself too.

Postgres has excellent full-text search capabilities. pgvector has excellent vector search capabilities. Combining them in a meaningful way? That’s on you.

You need to:

• Decide how to weight vector similarity vs. text relevance
• Normalize scores from two different scoring systems
• Tune the balance for your use case
• Probably implement Reciprocal Rank Fusion or something similar

Again, not impossible. Just another thing that many dedicated vector databases provide out of the box.

pgvectorscale (it doesn’t solve everything)

Timescale has released pgvectorscale, which addresses some of these issues. It adds:

• StreamingDiskANN, a new search backend that’s more memory-efficient
• Better support for incremental index builds
• Improved filtering performance

This is great! It’s also an admission that pgvector out of the box isn’t sufficient for production use cases.

pgvectorscale is still relatively new, and adopting it means adding another dependency, another extension, another thing to manage and upgrade. For some teams, that’s fine. For others, it’s just more evidence that maybe the “keep it simple, use Postgres” argument isn’t as simple as it seemed.

Oh, and if you’re running on RDS, pgvectorscale isn’t available. AWS doesn’t support it. So enjoy managing your own Postgres instance if you want these improvements, or just… keep dealing with the limitations of vanilla pgvector.

The “just use Postgres” simplicity keeps getting simpler.

Just use a real vector database

I get the appeal of pgvector. Consolidating your stack is good. Reducing operational complexity is good. Not having to manage another database is good.

But here’s what I’ve learned: for most teams, especially small teams, dedicated vector databases are actually simpler.

What you actually get

With a managed vector database (Pinecone, Weaviate, Turbopuffer, etc.), you typically get:

• Intelligent query planning for filtered searches
• Hybrid search built in
• Real-time indexing without memory spikes
• Horizontal scaling without complexity
• Monitoring and observability designed for vector workloads

It’s probably cheaper than you think

Yes, it’s another service to pay for. But compare:

• The cost of a managed vector database for your workload
• vs. the cost of over-provisioning your Postgres instance to handle index builds
• vs. the engineering time to tune queries and manage index rebuilds
• vs. the opportunity cost of not building features because you’re fighting your database

Turbopuffer starts at $64 month with generous limits.

For a lot of teams, the managed service is actually cheaper.

What I wish someone had told me

pgvector is an impressive piece of technology. It brings vector search to Postgres in a way that’s technically sound and genuinely useful for many applications.

But it’s not a panacea. Understand the tradeoffs.

If you’re building a production vector search system:

1. Index management is hard. Rebuilds are memory-intensive, time-consuming, and disruptive. Plan for this from day one.

2. Query planning matters. Filtered vector search is a different beast than traditional queries, and Postgres’s planner wasn’t built for this.

3. Real-time indexing has costs. Either in memory, in search quality, or in engineering time to manage it.

4. The blog posts are lying to you (by omission). They’re showing you the happy path and ignoring the operational reality.

5. Managed offerings exist for a reason. There’s a reason that Pinecone, Weaviate, Qdrant, and others exist and are thriving. Vector search at scale has unique challenges that general-purpose databases weren’t designed to handle.

The question isn’t “should I use pgvector?” It’s “am I willing to take on the operational complexity of running vector search in Postgres?”

For some teams, the answer is yes. You have database expertise, you need the tight integration, you’re willing to invest the time.

For many teams—maybe most teams—the answer is probably no. Use a tool designed for the job. Your future self will thank you.

Related reading

• RAG: From Context Injection to Knowledge Integration — why you’re storing these embeddings in the first place, and where the pattern breaks down.
• CheeseGPT — a small end-to-end RAG build, before any of the production concerns above apply.
• A Production Framework for LLM Feature Evaluation — deciding which LLM features justify this kind of infrastructure at all.

After several years of integrating LLMs into production systems, I’ve observed a consistent pattern: the features that deliver real value rarely align with what gets attention at conferences. While the industry focuses on AGI and emergent behaviors, the mundane applications—data extraction, classification, controlled generation—are quietly transforming how we build software.

This post presents a framework I’ve developed for evaluating LLM features based on what actually ships and scales. It’s deliberately narrow in scope, focusing on patterns that have proven reliable across multiple deployments rather than exploring the theoretical boundaries of what’s possible.

The Three Categories That Actually Work

Through trial, error, and more error, I’ve found that LLMs consistently excel in three specific areas. When I’m evaluating a potential AI feature, I ask: “Does this clearly fit into one of these categories?” If not, it’s probably not worth pursuing (yet).

1. Extracting Structured Data from Unstructured Inputs

This is the unsexy workhorse of AI features. Think of it as having an intelligent data entry assistant who never gets tired of parsing messy inputs.

What makes this valuable:

• Humans hate data entry
• Traditional parsing is brittle and breaks with slight format changes
• LLMs can handle ambiguity and variations gracefully

Real examples I’ve built:

• PDF to JSON converter: Taking uploaded forms (PDFs, images, even handwritten docs) and extracting structured data. What used to require complex OCR pipelines and regex nightmares now works with a simple prompt.
• API response mapper: Taking inconsistent third-party API responses and mapping them to your internal data model. Every integration engineer’s nightmare—different field names, nested structures that change randomly, optional fields that are sometimes null and sometimes missing entirely.
• Customer feedback analyzer: Extracting actionable insights from the stream of unstructured feedback across emails, Slack, support tickets. Automatically pulling out feature requests, bug reports, severity, and sentiment. What used to be a PM’s full-time job.

The key insight here is that LLMs excel at handling structural variance and ambiguity—the exact things that make traditional parsers brittle. A single well-crafted prompt can replace hundreds of lines of mapping logic, regex patterns, and edge case handling. The model’s ability to understand intent rather than just pattern match is what makes this category so powerful.

Production considerations: For high-volume extraction from standardized formats, purpose-built services like Reducto offer better economics and reliability than raw LLM calls. These platforms have already solved for edge cases around OCR quality, table extraction, and format variations. The build-vs-buy calculation here typically favors buying unless you have unique requirements or scale that justifies the engineering investment.

2. Content Generation and Summarization

This is probably what most people think of when they hear “AI features,” but the key is being specific about the use case.

What makes this valuable:

• Reduces cognitive load on users
• Provides consistent quality and tone
• Can process and synthesize large amounts of information quickly

Real examples I’ve built:

• Smart report generation: Taking raw data and generating human-readable reports with insights and recommendations.
• Meeting summarizer: Processing transcripts to extract key decisions, action items, and important discussions.
• Documentation assistant: Generating first drafts of technical documentation from code comments and README files.

The critical lesson here is that unconstrained generation is rarely what you want in production. Effective generation features require explicit boundaries: output structure, length constraints, tone guidelines, and forbidden topics. The challenge isn’t getting the model to generate—it’s getting it to generate within your specific constraints reliably.

This is where prompt engineering transitions from art to engineering: defining schemas, enforcing structural requirements, and building validation layers. The most successful generation features I’ve seen treat the LLM as one component in a larger pipeline, not a magic box.

3. Categorization and Classification

This is where LLMs really shine compared to traditional ML. What used to require thousands of labeled examples and complex training pipelines can now be done with a well-crafted prompt.

What makes this valuable:

• No need for labeled training data
• Can handle edge cases and ambiguity
• Easy to adjust categories without retraining

The architectural advantage here is profound: you’re essentially defining classifiers declaratively rather than imperatively. No training data, no model selection, no hyperparameter tuning—just clear descriptions of your categories. The model’s pre-trained understanding of language and context does the heavy lifting.

This fundamentally changes the iteration cycle. Adding a new category or adjusting definitions happens in minutes, not weeks. The trade-off is less fine-grained control over the decision boundary, but for most business applications, this is a feature, not a bug.

That said, “just use an LLM” isn’t automatically the right call for classification. I ran 32 experiments pitting small LLMs against a fine-tuned BERT encoder — the results are a useful counterweight to this section when latency and cost matter.

Scaling considerations: Production deployments require:

• Structured output guarantees: Libraries like Pydantic AI and Outlines enforce schema compliance at the token generation level, eliminating post-processing failures.
• Prompt optimization: DSPy and similar frameworks apply optimization techniques to prompt engineering, treating it as a learnable parameter rather than a manual craft.
• Evals, Observability, and Error Analysis: This could and will likely eventually be its own post

The Anti-Patterns: What Doesn’t Work

Let me save you some pain by sharing what consistently fails:

1. Trying to Replace Domain Expertise

LLMs are great at general knowledge but terrible at specialized domains without extensive context. If you need deep expertise, you still need experts.

2. Real-time, High-frequency Operations

Sub-100ms response times and high-frequency calls remain outside the practical envelope for LLM applications. The latency floor of current models, even with optimizations like speculative decoding, makes them unsuitable for hot-path operations.

3. Anything Requiring Perfect Accuracy

LLMs are probabilistic. If you need 100% accuracy (financial calculations, legal compliance, etc.), use traditional code.

A Practical Evaluation Framework

When someone comes to me with an AI feature idea, here’s my checklist:

Implementation Strategy

For teams evaluating their first LLM feature, I recommend starting with categorization. The reasoning is purely pragmatic: it has the clearest evaluation metrics, the most forgiving failure modes, and provides immediate value. You can validate the approach with a small dataset and scale incrementally.

The implementation complexity is also minimal—you’re essentially building a discriminator rather than a generator, which sidesteps many of the challenges around hallucination, output formatting, and content safety. Most importantly, when classification confidence is low, you can gracefully fall back to human review without breaking the user experience.

The Reality of Production AI

The gap between AI demos and production systems remains vast. The features that succeed in production share a common trait: they augment existing workflows rather than attempting to replace them entirely. They handle the tedious, error-prone tasks that humans perform inconsistently, freeing cognitive capacity for higher-value work.

This isn’t a limitation—it’s the current sweet spot for LLM applications. The technology excels at tasks that are simultaneously too complex for traditional automation but too mundane to justify human attention. Understanding this paradox is key to building AI features that actually ship.

Related reading

• Beating BERT? — 32 experiments on whether small LLMs actually beat a fine-tuned encoder at classification.
• RAG: From Context Injection to Knowledge Integration — the architectural limits of the most common LLM feature pattern.
• The Case Against pgvector — what the retrieval layer costs you once it’s in production.

I am 100% the guy in this panel right now.

But 4o image generation is insane.

I’ve been working in the LLM space since before ChatGPT shifted everything. I’ve closely followed the progress. I test every new release. I tell my friends that every AI app they send me is slop. I am not easily impressed.

But this feels like another ChatGPT moment. This isn’t just better distribution (is hiding your state-of-the-art model in a Discord chat behind /commands really the best way to get people to use it?).

This feels foundational. It’s not just a better diffusion model—it’s actual reasoning in pixel space.

I’ve been on the fence about whether AGI (whatever that even means) is possible. Can we actually bottle intelligence into an electric rock? But it doesn’t take much napkin math to pencil this out a few years. (True believers might ask where I’ve been, but rest easy brethren—I am yours now.)

It brings to mind this 100% real needlepoint of an Ilya Sutskever quote.

Trying to game out the second- and third-order effects of an image generation model feels strange, even dumb. Infinite Ghibli? What are you worried about? Ghibli gonna take all the jobs?

I think this tweet (shared with me by a friend this morning) just about sums it up.

But to that, I say:

Below is a series I’ve been working on to try and demonstrate this phenomenon I’m experiencing. My fiancée and I got engaged last October, and we captured an amazing photo (maybe my favorite picture ever). So I’ve been trying to recreate it in every style possible.

The consistency of this model is incredible–and the content filters are tuned to low right now. I won’t be surprised if by the time you’re reading this, most of these styles will be blocked (already seems to be happening :/ )

Our original photo:

Lego: Claymation: Sesame Street: Scooby-Doo: Neon Sign: Tim Burton: Hey Arnold: Victorian Botanical Print: Wes Anderson: Pixar: Vintage Comic: Peanuts: “Yellow Submarine Family”: Construction Paper: Architectural Blueprint: Medieval Manuscript: Street Art Stencil: Pixelated Video Game: 1960s Style Cartoon: Stop Motion:

And finally, Ghibli:

Other models could already do this!

No, they couldn’t.

Retrieval-Augmented Generation (RAG) has rapidly become a cornerstone in the practical application of Large Language Models (LLMs). Its promise is compelling: to expand LLMs beyond their training data by connecting them to external knowledge sources – from enterprise databases and real-time data streams to proprietary knowledge bases. The allure of RAG lies in its apparent simplicity – augment the LLM’s input context with retrieved information, and witness enhanced output quality. However, beneath this layer of simplicity lies a more complex reality–its a bit of a hack. RAG only works because LLMs are generally robust. The more you think on it, the more it becomes clear it shouldn’t really work, and should serve only as a stepping stone to a new paradigm.

Generation vs. Retrieval

At their core, LLMs are generative models that produce text by navigating through a high-dimensional latent space. During pre-training on large datasets, these models learn to map language into this space, capturing relationships between words, phrases, and concepts. Text generation isn’t a simple lookup process - it’s a sequential operation where the model predicts each token based on both the previous context and its learned representations.

RAG changes this core process significantly. Rather than relying only on the model’s learned representations, RAG injects external information directly into the context window alongside the user’s query. While this works well in practice, it raises important questions about the theoretical and architectural implications:

1. Impact on Generation Quality: How does inserting external information affect the model’s learned generation process? Does mixing training-derived and retrieved information create inconsistencies in the model’s outputs?

2. Information Integration: Can the model effectively combine information from different sources during generation? Or is it simply stitching together pieces without truly understanding how they relate?

3. Architectural Fitness: Are transformer architectures and their training objectives actually suited for combining retrieved information with generation? Or are we forcing an approach that doesn’t align with how these models were designed to work?

Real-World Limitations

These theoretical concerns manifest in several practical ways:

1. Context Integration Problems

Current RAG implementations often struggle with:

• Abrupt transitions between retrieved content and generated text
• Inconsistent voice and style when mixing sources
• Difficulty maintaining coherent reasoning across retrieved facts
• Limited ability to synthesize information from multiple sources

2. Attention Mechanism Overload

The transformer’s attention mechanism faces significant challenges:

• Managing attention across disconnected chunks of information
• Balancing focus between query, retrieved content, and generated text
• Handling potentially contradictory information from different sources
• Maintaining coherence when dealing with multiple retrieved documents

3. Knowledge Conflicts

RAG systems often struggle to resolve conflicts between:

• The model’s pretrained knowledge
• Retrieved information
• Different retrieved sources
• User queries and retrieved content

The Path Forward: Beyond Basic RAG

Recent research and development suggest several promising directions for addressing these limitations:

1. Improved Knowledge Integration

Future systems might:

• Process retrieved information before injection
• Maintain explicit source tracking throughout generation
• Use structured knowledge representations
• Implement hierarchical attention mechanisms

2. Enhanced Source Handling

Advanced approaches could:

• Evaluate source reliability and relevance
• Resolve conflicts between sources
• Maintain provenance information
• Generate explicit citations and references

3. Architectural Innovations

New architectures might include:

• Dedicated pathways for retrieved information
• Specialized attention mechanisms for source integration
• Dynamic context window management
• Explicit fact-checking mechanisms

The Next Evolution: Anthropic’s Citations API

Anthropic’s Citations API represents a significant step beyond traditional RAG implementations. While the exact implementation details aren’t public, we can make informed speculations about its architectural innovations based on the capabilities it demonstrates.

Architectural Innovations

The Citations API likely goes beyond simple prompt engineering to include fundamental architectural changes:

1. Enhanced Context Processing

   • Specialized attention mechanisms for source document processing
   • Dedicated layers for maintaining source awareness throughout generation
   • Architectural separation between query processing and source document handling
   • Advanced chunking and document representation strategies

2. Citation-Aware Generation

   • Built-in tracking of source-claim relationships
   • Automatic detection of when citations are needed
   • Dynamic weighting of source relevance
   • Real-time fact verification against sources

3. Training Innovations

   • Custom loss functions for citation accuracy
   • Source fidelity metrics during training
   • Explicit training for source grounding
   • Specialized datasets for citation learning

Speculation on Implementation

The system likely employs several key mechanisms:

1. Dual-Stream Processing

   • Separate processing paths for user queries and source documents
   • Specialized attention heads for citation tracking
   • Fusion layers for combining information streams
   • Dynamic context management

2. Source Integration

   • Fine-grained document chunking
   • Semantic similarity tracking
   • Citation boundary detection
   • Provenance preservation

3. Training Approach

   • Multi-task training combining generation and citation
   • Custom datasets focused on source grounding
   • Citation-specific loss functions
   • Source fidelity metrics

Beyond Traditional RAG

The Citations API and similar emerging technologies point to a future where knowledge integration isn’t just an add-on but a core capability of language models. This evolution requires moving beyond simply stuffing context windows with retrieved documents toward architectures specifically designed for knowledge-aware generation.

The next generation of these systems will likely feature:

• Native citation capabilities
• Real-time fact verification
• Seamless source integration
• Dynamic knowledge updates
• Explicit handling of source conflicts

As we move forward, the goal isn’t to patch the limitations of current RAG systems but to fundamentally rethink how we combine language models with external knowledge. This might lead to entirely new architectures specifically designed for knowledge-enhanced generation, moving us beyond the current paradigm of context window injection toward truly integrated knowledge-aware AI systems.

Related reading

• CheeseGPT — a simple, end-to-end RAG system built from scratch, if you want the concrete version of everything above.
• The Case Against pgvector — the retrieval half of RAG has to run somewhere, and the storage layer has sharper edges in production than most tutorials admit.
• A Production Framework for LLM Feature Evaluation — which LLM features actually survive contact with production.

If you just want the mental model:

• async def defines a coroutine function. Calling it doesn’t run anything — it returns a coroutine object, which is a description of work, not the work itself.
• await pauses the current coroutine and hands control back to the event loop, which is then free to run something else until the awaited thing finishes.
• The event loop is a single-threaded scheduler. There is no parallelism in async Python. It’s very efficient turn-taking, and the turns only change at an await.
• Async helps when you’re I/O-bound — waiting on networks, APIs, databases, disks. It does nothing for CPU-bound work; that’s what multiprocessing is for.
• One blocking call stalls everything. A time.sleep() or a synchronous DB driver anywhere in a coroutine freezes the entire loop, not just that coroutine.

The rest of this post is about why each of those is true, and what’s happening underneath.

Why async Python exists

Modern applications — especially anything dealing with network requests or many concurrent operations — need to handle multiple things at once. For a long time, threading and multiprocessing were the go-to answers in Python. But the Global Interpreter Lock and the overhead of thread management limit how well those approaches work, particularly for I/O-bound workloads where the CPU is mostly sitting idle waiting on something external.

Async is the alternative. The surface syntax of async and await looks simple enough, but understanding what’s happening beneath it is what separates working async code from code that mysteriously runs no faster than the synchronous version. We’ll cover the event loop, coroutines, tasks, futures, context switching, and how async compares to threads and processes.

async def and await: the syntax

These two keywords are the building blocks for writing asynchronous code in Python.

What async def actually does

In Python, you declare a function as asynchronous by using the async def syntax instead of the regular def. This seemingly small change has profound implications. An async def function, also known as a coroutine function, doesn’t execute like a regular synchronous function. Instead, it returns a coroutine object.

Think of a coroutine object as a promise of work to be done later. It’s not the work itself, but rather a representation of that work, ready to be executed when the time is right.

 1import asyncio
 2
 3async def hello_async():
 4    print("Hello from async function!")
 5    await asyncio.sleep(1) # Simulate some async operation (like network I/O)
 6    print("Async function finished.")
 7
 8# Calling the async function returns a coroutine object
 9coro = hello_async()
10print(f"Coroutine object: {coro}")
11
12# To actually run the coroutine, we need to use an event loop
13asyncio.run(coro)

In this example, hello_async is an asynchronous function. When we call hello_async(), it doesn’t immediately print “Hello from async function!”. Instead, it creates and returns a coroutine object. To actually execute the code within the coroutine, we need to use asyncio.run(), which sets up and runs an event loop (more on event loops shortly) to manage the execution of our coroutine.

What await actually does

The await keyword is the other half of the async duo. It can only be used inside async def functions. await is the point where an asynchronous function can pause its execution and yield control back to the event loop. This is the crucial mechanism that enables concurrency in async Python without relying on threads.

When you await something, you are essentially saying: “I need to wait for this asynchronous operation to complete. While I’m waiting, I’m going to yield control back to the event loop so it can work on other tasks. Once this operation is done, please resume my execution from right here.”

In our hello_async example, await asyncio.sleep(1) is a simulated asynchronous operation that represents waiting for 1 second. During this second, the hello_async coroutine pauses, and the event loop is free to execute other coroutines or handle other events. Once the sleep duration is over, the event loop resumes the hello_async coroutine from the line after the await statement, and it prints “Async function finished.”

Important Note: You can only await objects that are awaitable. In practice, this usually means you’re awaiting other coroutines, Future objects (which we’ll discuss later), or objects that have implemented the __await__ special method. Standard synchronous functions are not awaitable.

A complete async Python example

Here’s the thing the syntax alone doesn’t show you: await on its own doesn’t buy you any concurrency. If you await three things one after another, you wait for all three in sequence — exactly like synchronous code, just with extra keywords.

 1import asyncio
 2import time
 3
 4async def fetch(name, seconds):
 5    print(f"{name}: starting")
 6    await asyncio.sleep(seconds)   # stands in for a network call
 7    print(f"{name}: done")
 8    return f"{name} result"
 9
10async def sequential():
11    start = time.perf_counter()
12    await fetch("A", 1)
13    await fetch("B", 1)
14    await fetch("C", 1)
15    print(f"sequential took {time.perf_counter() - start:.1f}s")
16
17async def concurrent():
18    start = time.perf_counter()
19    results = await asyncio.gather(
20        fetch("A", 1),
21        fetch("B", 1),
22        fetch("C", 1),
23    )
24    print(f"concurrent took {time.perf_counter() - start:.1f}s")
25    return results
26
27asyncio.run(sequential())   # sequential took 3.0s
28asyncio.run(concurrent())   # concurrent took 1.0s

Same coroutines, same await, 3x difference. The concurrency comes from asyncio.gather() handing the event loop three coroutines at once, so it can start all three and then sit on the await while every one of them is in flight. In the sequential version the loop only ever knows about one at a time.

This is the single most common mistake in async Python: writing await in a loop and wondering why nothing got faster.

1# Slow — one at a time, despite being "async"
2results = []
3for url in urls:
4    results.append(await fetch_url(url))
5
6# Fast — all in flight at once
7results = await asyncio.gather(*(fetch_url(url) for url in urls))

One more trap worth knowing early. Because the event loop is single-threaded, a blocking call inside a coroutine doesn’t just block that coroutine, it blocks everything:

1async def bad():
2    time.sleep(1)          # blocks the entire event loop for a full second
3    
4async def good():
5    await asyncio.sleep(1) # yields control; other coroutines run

The same applies to any synchronous library — a non-async database driver, requests instead of aiohttp, a heavy CPU computation. If you must call blocking code, push it off the loop with asyncio.to_thread().

The asyncio event loop

At the heart of async Python lies the event loop. Think of the event loop as the central conductor of an orchestra. It’s responsible for managing and scheduling the execution of all your asynchronous tasks. It’s a single-threaded loop that constantly monitors for events and dispatches tasks to be executed when those events occur.

How the event loop works

1. Task Queue: The event loop maintains a queue of tasks (usually coroutines wrapped in Task objects) that are ready to be executed or resumed.
2. Event Monitoring: The event loop also monitors for various events, such as network sockets becoming ready for reading or writing, timers expiring, or file operations completing. It typically uses efficient system calls like select, poll, or epoll (depending on the operating system) to monitor these events without blocking.
3. Task Execution and Resumption: When an event occurs that makes a task ready to proceed (e.g., data is available on a socket that a task is waiting to read from), the event loop picks up that task from the queue and executes it until it encounters an await statement.
4. Yielding Control with await: When a coroutine reaches an await statement, it effectively tells the event loop, “I need to wait for this operation. Please pause me and let someone else run.” The event loop then takes control and looks for other tasks in the queue that are ready to run.
5. Resuming Execution: Once the awaited operation completes (e.g., the network request returns, the timer expires), the event loop is notified. It then puts the paused coroutine back into the task queue, ready to be resumed at the point where it left off.
6. Looping Continuously: The event loop continues this process of monitoring events, executing tasks, and pausing/resuming coroutines in a loop until there are no more tasks to run or the program is explicitly stopped.

asyncio, uvloop, and other event loop implementations

Python’s standard library provides the asyncio module, which includes a built-in event loop implementation. This default event loop is written in Python and is generally sufficient for many use cases.

However, for performance-critical applications, especially those dealing with high-performance networking, you might consider using alternative event loop implementations. One popular option is uvloop.

uvloop: uvloop is a blazing-fast, drop-in replacement for asyncio’s event loop. It’s written in Cython and built on top of libuv, the same high-performance library that powers Node.js. uvloop is significantly faster than the default asyncio event loop, especially for network I/O.

To use uvloop, you typically need to install it separately (pip install uvloop) and then set it as the event loop policy for asyncio when your application starts:

 1import asyncio
 2import uvloop
 3
 4async def main():
 5    print("Running with uvloop!")
 6    await asyncio.sleep(1)
 7    print("Done.")
 8
 9if __name__ == "__main__":
10    uvloop.install() # Set uvloop as the event loop policy
11    asyncio.run(main())

Other event loop implementations exist, but asyncio and uvloop are the most commonly used in Python async programming. Choosing between them often depends on the performance requirements of your application. For most general async tasks, asyncio’s default loop is perfectly adequate. For high-load network applications, uvloop can provide a noticeable performance boost.

Coroutines, tasks, and futures

To truly understand async Python, it’s helpful to think about it in layers. We’ve already touched upon coroutines and the event loop. Let’s now delve into the roles of Tasks and Futures.

Coroutines

As we discussed earlier, coroutines are the asynchronous functions you define using async def. They represent units of asynchronous work. Coroutines themselves are not directly executed by the event loop. Instead, they need to be wrapped in something that the event loop can manage and schedule. This “something” is a Task.

Tasks: running coroutines concurrently

A Task in asyncio is essentially a wrapper around a coroutine that allows the event loop to schedule and manage its execution. When you want to run a coroutine concurrently within the event loop, you typically create a Task from it.

You can create a Task using asyncio.create_task():

 1import asyncio
 2
 3async def my_coroutine():
 4    print("Coroutine started")
 5    await asyncio.sleep(2)
 6    print("Coroutine finished")
 7    return "Coroutine result"
 8
 9async def main():
10    task1 = asyncio.create_task(my_coroutine()) # Create a Task from the coroutine
11    task2 = asyncio.create_task(my_coroutine()) # Create another Task
12
13    print("Tasks created, waiting for completion...")
14
15    result1 = await task1 # Await the completion of task1
16    result2 = await task2 # Await the completion of task2
17
18    print(f"Task 1 result: {result1}")
19    print(f"Task 2 result: {result2}")
20
21asyncio.run(main())

In this example, we create two Tasks, task1 and task2, from the same my_coroutine. asyncio.create_task() schedules these coroutines to be run by the event loop concurrently. When we await task1 and await task2, we are waiting for these Tasks to complete and retrieve their results.

Tasks are essential for managing the lifecycle of coroutines within the event loop. They provide methods to:

• Cancel a Task: task.cancel()
• Check if a Task is done: task.done()
• Get the result of a Task: task.result() (if done)
• Get exceptions raised during Task execution: task.exception() (if any)

Futures: results that haven’t arrived yet

A Future is an object that represents the eventual result of an asynchronous operation. It’s a placeholder for a value that might not be available yet. Tasks in asyncio are actually a subclass of Futures.

Futures are used extensively in async Python to represent the outcome of operations that are performed asynchronously, such as:

• Network I/O: Reading data from a socket, sending a request to a server.
• File I/O: Reading or writing to a file (in an async context).
• Concurrent computations: Tasks running in parallel (within the same event loop or in different threads/processes).

A Future object has a state that can be:

• Pending: The asynchronous operation is still in progress.
• Running: The operation is currently being executed.
• Done: The operation has completed successfully or with an exception.
• Cancelled: The operation has been cancelled.

You can interact with a Future to:

• Check if it’s done: future.done()
• Get the result: future.result() (blocks until done if pending, raises exception if an exception occurred)
• Get exceptions: future.exception() (returns exception if one occurred, otherwise None)
• Add callbacks: future.add_done_callback(callback_function) (run a function when the future is done)
• Cancel the future: future.cancel()

When you await a Task (or any awaitable Future-like object), you are essentially waiting for that Future to become “done” and then retrieving its result or handling any exceptions.

How coroutines pause and resume

Now let’s look at how coroutines are actually executed, and how context switching works in async Python.

Cooperative vs. preemptive multitasking

Async Python uses cooperative multitasking. This is in contrast to preemptive multitasking used by operating systems for threads and processes.

• Preemptive Multitasking (Threads/Processes): In preemptive multitasking, the operating system’s scheduler decides when to switch between threads or processes. It can interrupt a running thread/process at any time and switch to another, even if the running thread/process doesn’t explicitly yield control. This is typically based on time slices and priority levels.
• Cooperative Multitasking (Async Python): In cooperative multitasking, coroutines voluntarily yield control back to the event loop when they encounter an await statement. The event loop then decides which coroutine to run next. Context switching only happens at these explicit await points. A coroutine will continue to run until it reaches an await or completes.

This cooperative nature has important implications:

• No True Parallelism (within a single event loop): Within a single event loop running in a single thread, true parallelism is not achieved. Coroutines take turns running. If a coroutine doesn’t await frequently and performs long-running CPU-bound operations, it can block the event loop and prevent other coroutines from making progress.
• Responsiveness: Cooperative multitasking is excellent for I/O-bound tasks. While one coroutine is waiting for I/O, another can run, keeping the application responsive.
• Less Overhead: Context switching in cooperative multitasking is generally lighter than preemptive context switching between threads or processes. There’s less operating system overhead involved.
• Deterministic Behavior (mostly): Because context switching happens only at explicit await points, the execution flow of async code is often more predictable and easier to reason about compared to multithreaded code, which can have race conditions and unpredictable scheduling.

What happens at an await

When a coroutine reaches an await statement, several things happen:

1. await Expression: The expression after await (e.g., asyncio.sleep(1), another coroutine, a Future) must be awaitable.
2. Yielding Control: The coroutine effectively “pauses” its execution at the await point. It returns control back to the event loop.
3. Event Loop Takes Over: The event loop becomes active again. It looks at its task queue for other tasks that are ready to run.
4. Registering for Resumption: The coroutine, along with information about where it paused (the line after the await), is registered with the event loop as being “waiting” for the completion of the awaited operation.
5. Awaited Operation Proceeds: The awaited operation (e.g., network request, timer) proceeds asynchronously in the background (often managed by non-blocking system calls).
6. Event Notification: When the awaited operation is complete, the event loop receives a notification (e.g., socket becomes readable, timer expires).
7. Resuming the Coroutine: The event loop puts the paused coroutine back into the task queue, marked as ready to be resumed.
8. Coroutine Resumes: When the event loop gets around to executing this coroutine again, it resumes from the exact point where it was paused (right after the await statement). It now has access to the result of the awaited operation (if any).

This pause-and-resume mechanism is what allows asynchronous code to be written in a seemingly sequential style, even though it’s actually being executed in an interleaved and non-blocking manner.

Async vs. threads vs. multiprocessing in Python

It’s worth being clear about when async is the right choice and when threading or multiprocessing is the better fit. They solve different problems, and reaching for the wrong one is how you end up with async code that runs no faster than the version you replaced.

When to use async: I/O-bound work

Async Python excels in scenarios where your application is I/O-bound. This means that the primary bottleneck is waiting for external operations to complete, such as:

• Network requests: Fetching data from APIs, making HTTP requests, communicating with databases over a network.
• File I/O: Reading and writing to files (especially over a network file system).
• Waiting for user input: In GUI applications or interactive systems.

In these cases, the CPU is often idle while waiting for I/O operations. Async Python allows you to utilize this idle time by letting other coroutines run while one is waiting for I/O. It’s highly efficient for handling many concurrent I/O operations with minimal overhead.

Example: Web Server

A web server that handles many concurrent requests is a classic example where async Python shines. While one request is being processed (which often involves waiting for database queries, external API calls, etc.), the server can be handling other requests concurrently.

 1import asyncio
 2import aiohttp
 3from aiohttp import web
 4
 5async def fetch_data_from_api(api_url):
 6    async with aiohttp.ClientSession() as session:
 7        async with session.get(api_url) as response:
 8            return await response.json()
 9
10async def handler(request):
11    data = await fetch_data_from_api("https://api.example.com/data")
12    return web.json_response(data)
13
14async def main():
15    app = web.Application()
16    app.add_routes([web.get('/', handler)])
17    runner = web.AppRunner(app)
18    await runner.setup()
19    site = web.TCPSite(runner, 'localhost', 8080)
20    await site.start()
21    print("Server started at http://localhost:8080")
22    await asyncio.Event().wait() # Keep the server running indefinitely
23
24if __name__ == "__main__":
25    asyncio.run(main())

This example uses aiohttp, an async HTTP client and server library. The fetch_data_from_api coroutine performs an asynchronous HTTP request. The handler coroutine uses await fetch_data_from_api to fetch data without blocking the server. The server can handle many requests concurrently, making it highly scalable for I/O-bound web applications.

When to use threads

Threads, especially when used with Python’s threading module, are suitable for tasks that are more CPU-bound and can benefit from concurrency (even if not true parallelism due to the GIL).

• CPU-Bound Tasks: Tasks that spend most of their time performing computations on the CPU, rather than waiting for I/O. Examples include:

  • Image processing
  • Numerical computations
  • Data analysis
  • Cryptographic operations

• Concurrency (with GIL limitations): Python’s GIL (Global Interpreter Lock) prevents true parallelism for CPU-bound tasks in standard CPython threads. Only one thread can hold the Python interpreter lock at any given time. However, threads can still provide concurrency by releasing the GIL during I/O operations or certain blocking system calls. This can improve responsiveness even for CPU-bound tasks if they involve some I/O or blocking.

Example: CPU-Bound Computation (with threading for concurrency)

 1import threading
 2import time
 3
 4def cpu_bound_task(task_id):
 5    print(f"Task {task_id} started")
 6    start_time = time.time()
 7    result = 0
 8    for _ in range(10**7): # Simulate CPU-intensive computation
 9        result += 1
10    end_time = time.time()
11    duration = end_time - start_time
12    print(f"Task {task_id} finished in {duration:.4f} seconds, result: {result}")
13
14def main_threads():
15    threads = []
16    for i in range(4):
17        thread = threading.Thread(target=cpu_bound_task, args=(i,))
18        threads.append(thread)
19        thread.start()
20
21    for thread in threads:
22        thread.join() # Wait for all threads to complete
23
24if __name__ == "__main__":
25    main_threads()

In this example, cpu_bound_task simulates a CPU-intensive operation. We create multiple threads to run this task concurrently. While the GIL limits true parallelism for CPU-bound Python code, threads can still provide some concurrency and potential performance improvement, especially if the tasks involve some I/O or blocking operations. For purely CPU-bound tasks, however, the benefits might be limited by the GIL.

When to use processes: true parallelism

For truly CPU-bound and computationally intensive tasks that need true parallelism and to bypass the GIL limitations, multiprocessing using Python’s multiprocessing module is the way to go.

• True Parallelism: Multiprocessing creates separate processes, each with its own Python interpreter and memory space. Processes run in parallel on multiple CPU cores, achieving true parallelism for CPU-bound tasks.
• Bypassing the GIL: Each process has its own GIL, so the GIL limitation of threads is overcome.
• Higher Overhead: Process creation and inter-process communication have more overhead compared to threads or async tasks. Processes consume more system resources (memory, process management overhead).

Example: CPU-Bound Computation (with multiprocessing for parallelism)

 1import multiprocessing
 2import time
 3
 4def cpu_bound_task_process(task_id): # Same CPU-bound task as before
 5    print(f"Process {task_id} started")
 6    start_time = time.time()
 7    result = 0
 8    for _ in range(10**7):
 9        result += 1
10    end_time = time.time()
11    duration = end_time - start_time
12    print(f"Process {task_id} finished in {duration:.4f} seconds, result: {result}")
13
14def main_processes():
15    processes = []
16    for i in range(4):
17        process = multiprocessing.Process(target=cpu_bound_task_process, args=(i,))
18        processes.append(process)
19        process.start()
20
21    for process in processes:
22        process.join() # Wait for all processes to complete
23
24if __name__ == "__main__":
25    main_processes()

In this multiprocessing example, we create separate processes to run the same CPU-bound task. Because each process has its own interpreter and bypasses the GIL, we can achieve true parallelism and significantly speed up CPU-intensive computations on multi-core systems.

Choosing between async, threads, and processes

General Guidelines:

• I/O-Bound, High Concurrency: Async Python (asyncio) is often the best choice.
• CPU-Bound with some I/O, Responsiveness: Threads (threading) can be considered, but be mindful of GIL limitations for pure CPU-bound tasks.
• CPU-Bound, True Parallelism, Max Performance: Processes (multiprocessing) are essential, especially for computationally intensive tasks on multi-core machines.
• Hybrid Applications: You can combine async and multiprocessing. For example, use async for handling network I/O and multiprocessing for CPU-bound background tasks.

When to reach for async

Async Python is a good way to write concurrent code for I/O-bound applications, and understanding the machinery underneath — the event loop, coroutines, tasks, futures, cooperative multitasking — is what makes the difference between using it well and using it superstitiously.

It isn’t a silver bullet. It won’t speed up CPU-bound work, it doesn’t replace threading or multiprocessing, and a single blocking call will undo all of it. But for the case it’s designed for — lots of concurrent waiting on things outside your process — nothing else in Python comes close on overhead.

If you take three things from this post: async def returns a coroutine rather than running one, concurrency comes from gather and tasks rather than from await itself, and the loop is single-threaded so anything that blocks blocks everyone.

If you want to see these concepts in a real application, FastAPI is async Python end to end — I’ve written about integration testing FastAPI, where the event loop and async def shape how you have to structure your test fixtures.

Introduction

If you’re a backend developer working with FastAPI, you already know the framework excels at simplifying API development with features like async support and automated Swagger docs. However, when it comes to integration testing–particularly mocking external services like MongoDB, AWS S3, and third-party APIs—the waters can get murky. This post is your guide to navigating these complexities.

This isn’t a one-stop-shop for all things testing; the focus is squarely on integration testing within FastAPI.

Prerequisites: Familiarity with FastAPI, PyTest, MongoDB, and AWS S3 is assumed. If you’re new to any of these technologies, you may want to get up to speed before proceeding. FastAPI is built on async Python, and a lot of testing confusion traces back to not knowing what async def and the event loop are really doing — if that’s shaky ground, my deep dive into Python async programming covers it.

What is Integration Testing?

Integration testing involves combining individual units of code and testing them as a group. This type of testing aims to expose faults in the interactions between integrated units. In our context, we could also probably call these API tests (and you’ll see that’s how I name my test files). In the context of FastAPI, these units often involve your API endpoints, external databases like MongoDB, and other services such as AWS S3.

Why It’s Challenging

When compared to unit tests, integration tests in FastAPI present unique difficulties. These challenges mostly stem from the interaction with external dependencies. Mocking these dependencies is not always straightforward, and mistakes can lead to false positives or negatives, undermining the purpose of the tests.

Why It Matters

So, why focus on integration testing in FastAPI? Because it’s here that you validate that various services and databases work in harmony. By ensuring that these integrated units function as expected, you not only increase code reliability but also save debugging time in the long run.

Unit tests are great for testing individual units of code, but they don’t test the interactions between these units. I also find integration tests (especially this kind, which test our endpoints) are particularly useful before doing a large refactor–going from a v0 to v1 of an app, where unit tests might not make a ton of sense because you’re rewriting all of your ‘units’.

A Glimpse into Our Test App

Before we delve into the technicalities of mocking various elements, let’s take a quick look at the FastAPI application we’ll be using as a test subject. This app serves as a sandbox for our integration tests, built to showcase different aspects of FastAPI that commonly require mocking in a test environment.

Our sample application has a straightforward schema designed for demonstration purposes:

 1/login  # simple login
 2  POST: Login
 3
 4/users/me/  # example for mocking auth
 5  GET: Read Users Me
 6
 7/users/me/preferences  # example for mocking mongodB
 8  GET: Get Preferences
 9  POST: Save Preferences
10
11/users/me/profile_pic  # example for mocking s3
12  GET: Get Profile Picture
13  POST: Save Profile Picture
14
15/weather/{city}  # example for mocking external API
16  GET: Get Weather

Our tests live in tests/ where they are organized according to the application components they evaluate. We employ PyTest for test execution and adhere to its conventions for discovering tests and initial set up. This includes utilizing the conftest.py configuration file for shared hooks and fixtures.

Mocking Authentication in FastAPI

As we move through the article, we’ll see that FastAPI provides an easy way to do dependency injection, which is especially useful for our testing scenarios. Using dependency_overrides, we can inject our mocked functions into the FastAPI app, but remember, this only works for endpoints/functions using the Depends() syntax. You can read more about dependency injection in FastAPI here.

In our FastAPI application, we have endpoints that require authentication. During testing, we don’t want to use real authentication because it would require us to manage real user credentials and tokens. This would complicate our tests and potentially expose sensitive information. Therefore, we mock the authentication process.

Testing Our Login

Before we get to dependency injection or mocking out our services, let’s make sure our auth works. We have simple authentication defined in the project that checks if a user is in the database, if their password matches, and returns a token if they are. (I am not going to focus on how to do authentication, but you can check out the code in auth.py if you’re interested.)

For the purposes of this tutorial, our database is just a dictionary with one hard coded user.

(I’ll be using my name as the username in this tutorial in order to burn it into your brain)

1# File: app/db.py
2users_db: Dict[str, UserInDB] = {
3    "alexjacobs": UserInDB(
4        username="alexjacobs",
5        email="alex@example.com",
6        hashed_password=get_password_hash("secret"),
7    )
8}

and our login endpoint…

 1# File: app/main.py
 2@app.post("/login", response_model=Token)
 3async def login(_login: Login):
 4    user = authenticate_user(users_db, _login.username, _login.password)
 5    if not user:
 6        raise HTTPException(
 7            status_code=status.HTTP_401_UNAUTHORIZED,
 8            detail="Incorrect username or password",
 9            headers={"WWW-Authenticate": "Bearer"},
10        )
11    access_token = create_access_token(
12        data={"sub": user.username}
13    )
14    return {"jwt": access_token, "token_type": "bearer"}

and we have some simple tests just to verify that our login endpoint works as expected

 1# File: tests/test_auth_api.py
 2def test_login(client):
 3    response = client.post("/login", json={"username": "alexjacobs", "password": "secret"})
 4    assert response.status_code == 200
 5    assert "jwt" in response.json()
 6
 7
 8def test_login_fail(client):
 9    response = client.post("/login", json={"username": "alexjacobs", "password": "not_the_right_password"})
10    assert response.status_code == 401
11    assert response.json() == {'detail': 'Incorrect username or password'}

These initial tests don’t require mocking the authentication process because they are designed to test the endpoint’s basic functionality. They act as a foundational layer upon which we will build more complex test scenarios that require mocking.

Mocking Authentication in Our Test Client

Now, let’s say we have another endpoint that requires authentication. We’ll use the Depends function to inject our authentication function into our endpoint. Our /users/me endpoint returns information about the user logged in.

1# File: app/main.py
2@app.get("/users/me/", response_model=User)
3async def read_users_me(_user: TokenData = Depends(get_auth)):
4    user = users_db.get(_user.username)
5    if user is None:
6        raise HTTPException(status_code=404, detail="User not found")
7    return user

Authentication is handled by the get_auth function, which is injected into our endpoint using the Depends function.
Since we don’t want to use real authentication in our tests, we’re going to mock this function. I’m going to show multiple ways of doing this so you can choose the one that works best for you.

The most standard way of mocking this function is to use the dependency_overrides feature of FastAPI, before we initialize our TestClient.

 1# File: tests/conftest.py
 2@pytest.fixture
 3def client(mock_s3_bucket):
 4    # we patch auth within our client fixture
 5    from app.main import app
 6
 7    def mock_get_auth():
 8        return TokenData(username="alexjacobs")
 9
10    app.dependency_overrides[get_auth] = mock_get_auth
11    with TestClient(app) as test_client:
12        yield test_client
13
14    app.dependency_overrides.clear()

You’ll see we declare a function called mock_get_auth that returns a TokenData object. Then in line 10, app.dependency_overrides[get_auth] = mock_get_auth, we ‘inject’ our mock function into our app in place of the get_auth function.

This code is essentially replacing the get_auth function in our app with a mock function that returns the TokenData object hardcoded into the function we’re pathing with.

Then, in our test we pass in the client fixture, and we can see that our test passes.

1# File: tests/test_auth_api.py
2def test_get_me_patched_auth(client):
3    # auth works without real jwt because we patched it in the client fixture
4    response = client.get("/users/me", headers={"Authorization": "Bearer " + 'fake.jwt'})
5    assert response.status_code == 200
6    assert response.json() == {'username': 'alexjacobs', 'email': 'alex@example.com'}

(we really don’t even need to include auth headers, since we’re mocking the auth function, but I’m keeping it here for clarity)

Mocking Authentication Directly in your Tests

Now, what if we want to mock the auth function in a different way? We can do this by mocking the auth function directly in our test.

First, we’ll create a second client fixture that doesn’t patch the auth function.

1# File: tests/conftest.py
2@pytest.fixture
3def client_unpatched_auth():
4    # we don't patch auth for this client
5    from app.main import app
6
7    with TestClient(app) as test_client_2:
8        yield test_client_2

Now we’re going to write a test to verify that auth fails (since we’re not patching the auth function)

1# File: tests/test_auth_api.py
2def test_get_me_unpatched_auth(client_unpatched_auth):
3    # auth fails without real jwt because we're using the unpatched client fixture
4    response = client_unpatched_auth.get("/users/me", headers={"Authorization": "Bearer " + 'fake.jwt'})
5    assert response.status_code == 401
6    assert response.json() == {"detail": "Could not validate credentials"}

To make this work, there are two approaches. The first is to mock the auth function directly in the test

 1# File: tests/test_auth_api.py
 2def test_get_me_path_auth_in_fn(client_unpatched_auth):
 3    # auth works because we patch it in this test (not in the client fixture)
 4    def mock_get_auth():
 5        return TokenData(username="alexjacobs")
 6
 7    app.dependency_overrides[get_auth] = mock_get_auth
 8
 9    response = client_unpatched_auth.get("/users/me", headers={"Authorization": "Bearer " + 'fake.jwt'})
10    assert response.status_code == 200

All we’ve done here is moved the code that patches the auth function into the test itself. This could be useful if you wanted to mock the auth function in some tests, but not others (but only wanted to have one client fixture).

Next, we’ll do the same thing, but we’re going to use a new fixture and factory function to create our mock auth function. So first, we’ll create a new fixture,

1# File: tests/conftest.py
2@pytest.fixture
3def mock_get_auth_factory():
4    def mock_get_auth():
5        return TokenData(username="alexjacobs")
6
7    return mock_get_auth

and then we’ll use this fixture in our test

1# File: tests/test_auth_api.py
2def test_get_me_patch_auth_with_fixture(client_unpatched_auth, mock_get_auth_factory):
3    # auth works because we patch it with a fixture instead of in the client fixture
4    app.dependency_overrides[get_auth] = mock_get_auth_factory
5    response = client_unpatched_auth.get("/users/me", headers={"Authorization": "Bearer " + 'fake.jwt'})
6    assert response.status_code == 200

Notice that when we’re passing in the mock_get_auth_factory fixture, we’re passing in the function itself, not the return value of the function. This is because the fixture is a factory function that returns a function, so we need to pass in the function itself.

This may not seem very useful (and for mocking auth, it may not be), but there are plenty of scenarios when you may want to mock a function in some tests, but not others. This is a good way to do that.

Mocking External APIs

Our FastAPI application makes external API calls to fetch weather data. During testing, we don’t want to make real API calls because they can be slow and unreliable. Therefore, we mock the external API calls.

We mock the external API calls by patching the fetch_weather function in our FastAPI application. This function is responsible for making the actual API call and returning the weather data. We replace it with a mock function that always returns a predefined weather data. But, this function does use dependency injection (the Depends funciton) in our app, so we have to mock it in a different way.

First, let’s look at our endpoint and the function it calls.

1# File: app/main.py
2@app.get("/weather/{city}", response_model=WeatherResponse)
3async def get_weather(city: str, _user: TokenData = Depends(get_auth)):
4    weather = fetch_weather(city)
5    return WeatherResponse(city=city, weather=weather)

and our fetch_weather() function

1# File: app/helpers.py
2def fetch_weather(city_name):
3    url = f"http://wttr.in/{city_name}?format=%C+%t"
4    response = requests.get(url)
5    return response.text

Pretty simple… we are hitting an external API, though, so we need to mock this out in our tests.

We have three ways to mock the external API calls:

Mocking External APIs with unittest.mock.patch Directly in Our Test

Using unittest.mock.patch, we patch our fetch_weather function to return ‘sunny’. We then verify our endpoint returns the expected response (and verify that our mock function was called–this last part probably isn’t really necessary here, but is a good demonstration)

1# File: tests/test_weather_api.py
2def test_get_weather(client):
3    with patch('app.main.fetch_weather', return_value='sunny') as mock_fetch_weather:
4        response = client.get("/weather/atlanta")
5        assert response.status_code == 200
6        assert response.json() == {'city': 'atlanta', 'weather': 'sunny'}
7        mock_fetch_weather.assert_called_once_with('atlanta')

Creating a fixture that patches the function

Let’s suggest that we may want to test multiple endpoints that call the fetch_weather function. We could patch the function in each test, but that’s a lot of code duplication. Instead, we can create a fixture that patches the function. First, we’ll set up our fixture (this is another factor function). We’re using unittest’s Mock and Patch functions in order to replace the fetch_weather function with a mock function that returns ‘rainy’.

Note: As before, we’re using a factory function to create our mock function, so we need to pass in the function itself, not the return value of the function.

1# File: tests/conftest.py
2@pytest.fixture
3def mock_fetch_weather_factory():
4    def mock_fetch_weather(city):
5        return "rainy"
6
7    mock = Mock(side_effect=mock_fetch_weather)
8    with patch.object(app.main, 'fetch_weather', new=mock):
9        yield

In our test, we just need to again pass in the fixture as an argument.

1# File: tests/test_weather_api.py
2def test_get_weather1(client, mock_fetch_weather_factory):
3    response = client.get("/weather/atlanta")
4    assert response.status_code == 200
5    assert response.json() == {'city': 'atlanta', 'weather': 'rainy'}

Creating a parametrized fixture

The above example probably isn’t very useful in practice. In the real world, if we’re testing multiple endpoints that call the same function, we probably want to test different scenarios. For example, we may want to test that our endpoint returns the correct response for different weather conditions. We could do this by creating multiple fixtures that patch the function with different mock functions, but this is a lot of code duplication and basically defeats the purpose of using a fixture at all. Instead, we can use a parametrized fixture to pass in values to our fixture to make it behave differently depending on our test.

To do this, we’re going to create another fixture, but this one will take a parameter. It essentially wraps our previous factory function, but now we can pass in a parameter to the fixture.

 1# File: tests/conftest.py
 2@pytest.fixture
 3def mock_fetch_weather_parametrized(request):
 4    def mock_fetch_weather_factory(weather):
 5        def mock_fetch_weather(*args, **kwargs):
 6            return weather
 7
 8        return mock_fetch_weather
 9
10    mock_weather = request.param
11    with patch.object(app.main, 'fetch_weather', new=mock_fetch_weather_factory(mock_weather)):
12        yield

And then, when we call our test, we need to decorate it with pytest.mark.parametrize, and pass in the fixture as an argument. I’ve set up two tests here so we can really see how it works.

 1# File: tests/test_weather_api.py
 2@pytest.mark.parametrize('mock_fetch_weather_parametrized', ['cloudy'], indirect=True)
 3def test_get_weather_parameterized(client, mock_fetch_weather_parametrized):
 4    response = client.get("/weather/atlanta")
 5    assert response.status_code == 200
 6    assert response.json() == {'city': 'atlanta', 'weather': 'cloudy'}
 7
 8
 9@pytest.mark.parametrize('mock_fetch_weather_parametrized', ['tornados'], indirect=True)
10def test_get_weather_parameterized(client, mock_fetch_weather_parametrized):
11    response = client.get("/weather/atlanta")
12    assert response.status_code == 200
13    assert response.json() == {'city': 'atlanta', 'weather': 'tornados'}

The parameterized fixture is more complicated to set up, but it is incredibly useful in practice when you have to simulate different responses from external APIs.

Mocking MongoDB

Our FastAPI application interacts with MongoDB. During testing, we might not want to (or really be able to in some cases) hit a real database with real/mocked data. Instead, we acn mock MongoDB by using the mongomock library, which simulates a MongoDB client. (Note: In my experience, mongomock can be difficult to work with and some practice to get working, but were going to proceed with it for now)

Creating a Mock MongoDB Client

In our application, we employ context management for database interactions, utilizing Python’s with statement. This demands that the object used within the with statement must implement context management protocols, specifically __enter__ and __exit__ methods.

 1# File: app/db.py
 2@contextmanager
 3def get_mongodb():
 4    try:
 5        with MongoClient() as client:
 6            db = client["preferences"]
 7            yield db
 8    except Exception as e:
 9        print(f'Error: {e}')
10        raise

The mongomock library doesn’t implement these methods, so to align with these requirements, we define a custom MockMongoClient class to wrap our mongomock client. This class mimics the behavior of the actual MongoClient by implementing the __enter__ and __exit__ methods. This ensures compatibility with the existing code that expects a context-managed database client.

 1# File: tests/conftest.py
 2class MockMongoClient:
 3    def __init__(self, db):
 4        self.db = db
 5
 6    def __enter__(self):
 7        return self.db
 8
 9    def __exit__(self, *args):
10        pass

Using our MockMongoClient class, we can now create our mock database client. We’ll do this in a fixture so we can reuse it in multiple tests.

Creating a Mock MongoDB Client Fixtures

I’m going to initialize two separate fixtures here, one with an empty database and one with some data initialized. (In theory, we should be able to set the scope of the fixture to session or module and just initialize the database once, add data through other tests, and then use that data for testing later, but I haven’t been able to get that working with mongomock. The project seems to have been designed more with unit tests in mind and is not a complete implementation or drop in replacement. If you know how to make this work, please let me know!)

1# File: tests/conftest.py
2@pytest.fixture
3def mock_mongodb():
4    def mock_get_mongodb():
5        mock_client = MongoClient()
6        return MockMongoClient(mock_client.db)
7
8    return mock_get_mongodb

and initialize some data…

1# File: tests/conftest.py
2@pytest.fixture
3def mock_mongodb_initialized():
4    def mock_get_mongodb():
5        mock_client = MongoClient()
6        mock_client.db.preferences.insert_one({"username": "alexjacobs", "city": "Berlin"})
7        return MockMongoClient(mock_client.db)
8
9    return mock_get_mongodb

And now we can write our tests. We’ll pass in our fixture and again use the dependency_overrides feature to inject our mock database client into our app (overriding the get_mongodb function)

Writing Tests Integration Tests Using Our Mocked MongoDB Client

We’ve got three simple tests here.
First, we test that our endpoint correctly returns a 404 if no user preferences exist (we use our non-initialized mock_mongodb fixture for this)

Next, we use our initialized mock_mongodb fixture to test that our endpoint correctly returns the user preferences.

Finally, we test that we can save user preferences (we use our non-initialized mock_mongodb fixture for this, but it should work with either)

 1# File: tests/user_preference_api.py
 2from app.main import app
 3from app.fake_db import get_mongodb
 4
 5
 6def test_get_user_preferences_404(client, mock_mongodb):
 7    # should return 404 since no user preferences exist
 8    app.dependency_overrides[get_mongodb] = mock_mongodb
 9    response = client.get("/users/me/preferences")
10    assert response.status_code == 404
11    assert response.json() == {'detail': 'Preferences not found'}
12
13
14def test_get_user_preferences_initialized(client, mock_mongodb_initialized):
15    # we're using our mock_mongodb_initialized fixture here which has our user preferences already set
16    app.dependency_overrides[get_mongodb] = mock_mongodb_initialized
17    response = client.get("/users/me/preferences")
18    assert response.status_code == 200
19    assert response.json() == {'city': 'Berlin'}
20
21
22def test_post_user_preferences(client, mock_mongodb):
23    app.dependency_overrides[get_mongodb] = mock_mongodb
24    response = client.post("/users/me/preferences", json={"city": "Berlin"})
25    assert response.status_code == 200
26    assert response.json() == {"detail": "success"}

And that’s it! Pretty simple to write the actual tests once we have our mocked database client set up correctly.

Mocking AWS S3

Our FastAPI application interacts with AWS S3 for storing and retrieving user profile pictures. During testing, we don’t want to use a real S3 bucket because it would require us to manage real AWS resources. This would complicate our tests and potentially incur costs. Therefore, we mock the S3 bucket. (The same patterns used here could be applied to other AWS services)

Creating the Mock S3 Fixture

We mock the S3 bucket by using the mock_s3 decorator from the moto library, which simulates an S3 bucket. This is done in a fixture so we can reuse it in multiple tests.

(moto is a great library for mocking AWS services, and you’ll see we are able to set our fixture scope to session, allowing us to maintain bucket state across multiple tests)

 1# File: tests/conftest.py
 2@pytest.fixture(scope="session")
 3def mock_s3_bucket():
 4    with mock_s3():
 5        conn = boto3.resource('s3', region_name='us-east-1')
 6        conn.create_bucket(Bucket=bucket)
 7        # we could upload a test file(s) here if we wanted to
 8        # s3_client = boto3.client('s3', region_name='us-east-1')
 9        # s3_client.upload_file('tests/assets/duck.png', bucket, 'profile_pics/alexjacobs.png') 
10        yield

Integrating with Test Client

This looks similar to our MockMongo fixtures, but this time, rather than pass our fixture into our tests, we pass it into the client fixture.

 1# File: tests/conftest.py
 2@pytest.fixture
 3def client(mock_s3_bucket):
 4    # we patch auth within our client fixture
 5    from app.main import app
 6
 7    def mock_get_auth():
 8        return TokenData(username="alexjacobs")
 9
10    app.dependency_overrides[get_auth] = mock_get_auth
11    with TestClient(app) as test_client:
12        yield test_client
13
14    app.dependency_overrides.clear()

Writing the Tests

Now we can write our tests. We’ll start by testing that we get the right message if no profile picture exists for the user

1# File: tests / test_user_preference_api.py
2def test_get_user_profile_pic_404(client):
3    response = client.get("/users/me/profile_pic")
4    assert response.status_code == 404
5    assert response.json() == {'detail': 'No profile pic found'}

Now we’ll test adding one…

1# File: tests/test_user_preference_api.py
2def test_set_user_profile_pic(client):
3    with open('tests/assets/duck.png', 'rb') as f:
4        response = client.post("/users/me/profile_pic", files={"picture": ("duck.png", f, "image/png")})
5
6    assert response.status_code == 200
7    assert response.json() == {'detail': 'success'}

Finally, we can test getting one. (notice how this is using the same file we uploaded in the previous test, this is due to the fact that we’re using the session scope for our mock_s3_bucket fixture)

1# File: tests/test_user_preference_api.py
2def test_get_user_profile_pic(client):
3    response = client.get("/users/me/profile_pic")
4    assert response.status_code == 200
5    with open('tests/assets/duck.png', 'rb') as f:
6        original_image_data = f.read()
7        original_base64 = base64.b64encode(original_image_data).decode('utf-8')
8
9    assert response.json()['image'] == original_base64

Wrapping Up

This was a fairly deep dive. We’ve unraveled the intricacies of integration testing in FastAPI, which is not without its challenges when it comes to mocking external dependencies. We’ve gone through a variety of techniques to mock authentication, from simple dependency_overrides to more advanced fixture-based strategies. We’ve also tackled how to mock external APIs using Python’s unittest.mock.patch and pytest’s parametrized fixtures.

When it comes to databases, MongoDB adds another layer of complexity. We’ve seen how Mongomock can be a useful tool, albeit with its own set of limitations. We crafted custom mock MongoDB clients and fixtures to ease this pain point. As for AWS S3, the Moto library proved to be a robust tool, enabling us to mock S3 buckets effectively, even allowing state persistence across tests.

The aim has been to arm you with a set of tools and strategies for your FastAPI testing arsenal. Whether it’s a simple authentication mock or a more complex external service, you should now be equipped to tackle these head-on. Happy testing.

Happy testing.

This toy project was originally created for a guest lecture to a Data Science 101 course (and its quality may reflect that :) This post extends that lecture, designed to provide a high-level understanding and example of Retrieval Augmented Generation (RAG).

We’ll go through the steps of creating a RAG based LLM system, explaining what we’re doing along the way, and why. For the argument about why RAG works at all — and where it starts to break down architecturally — see RAG: From Context Injection to Knowledge Integration.

You can follow along with the slides and code here

The CheeseGPT System

CheeseGPT combines Large Language Models (LLMs) with the advanced capabilities of Retrieval-Augmented Generation (RAG). At its core, CheeseGPT uses OpenAI’s GPT-4 model for natural language processing. This model serves as the backbone for generating human-like text responses. However, what sets CheeseGPT apart is its integration with Langchain and a Redis database containing all of the information on Wikipedia relating to cheese.

When a user asks a question, the system utilizes RAG to retrieve the most relevant information/documents from its vector database, and then includes those in its message to the LLM. This allows the LLM to have specific and up-to-date information to use, extending from the data that it was trained on.

The image below, flow from right to left (steps 1-5) shows the high level design of this. The user’s query is passed into our embedding model. We do a similarity search against our database to retrieve the most relevant documents to our users question. And then these are included in context passed to our LLM.

https://www.anyscale.com/blog/a-comprehensive-guide-for-building-rag-based-llm-applications-part-1

Below, we’ll outline the steps to building this system.

NOTE: this is an example, and probably doesn’t make a ton of sense as a useful system. (For one, we’re getting our data from wikipedia, which is already contained within the training data of GPT-4) This is meant to be a high level example that can show how a RAG based system can work, and to show what the possibilities are when integrating external data with LLMs (proprietary data, industry specific technical docs, etc.)

Data Collection and Processing

As with most projects, getting and munging your data is one of the most time consuming yet crucial elements. For our CheeseGPT example, this involved scraping Wikipedia for cheese-related articles, generating embeddings, and storing them in a Redis database. Below, I’ll outline these steps with code snippets for clarity.

Scraping Wikipedia

We start by extracting content from Wikipedia. We made a recursive function get_page_content to fetch pages related to cheese, including summaries and sections. (Note: This function could definitely be improved.)

 1
 2def get_page_content(page_title, depth, max_depth):
 3    wiki_wiki = wikipediaapi.Wikipedia('MyCheeseRAGApp/1.0 (myemail@example.com)')
 4    if depth > max_depth or page_title in visited_pages:
 5        return [], []
 6    visited_pages.add(page_title)
 7
 8    if len(visited_pages) % 10 == 0:
 9        logger.info(f"Visited pages: {len(visited_pages)}")
10        logger.info(f"Fetching page '{page_title}' (depth={depth})")
11
12    try:
13        page = wiki_wiki.page(page_title)
14        if not page.exists():
15            return [], []
16
17        texts, metadata = [], []
18        if len(texts) % 100:
19            logger.info(f"Texts length {len(texts)}")
20
21        # Add page summary
22        texts.append(page.summary)
23        metadata.append({'title': page_title, 'section': 'Summary', 'url': page.fullurl})
24
25        # Add sections
26        for section in page.sections:
27            texts.append(section.text)
28            metadata.append({'title': page_title, 'section': section.title, 'url': page.fullurl})
29
30        # Recursive fetching for links
31        if depth < max_depth:
32            for link_title in page.links:
33                link_texts, link_metadata = get_page_content(link_title, depth + 1, max_depth)
34                if len(link_texts) > 0:
35                    texts.extend(link_texts)
36                    metadata.extend(link_metadata)
37
38        return texts, metadata
39    except Exception as e:
40        logger.error(f"Error fetching page '{page_title}': {e}")
41        return [], []

This is a very greedy (and lazy) approach. We don’t discriminate at all, and we end up with a ton of noise (things not related to cheese at all), but for our purposes of example, it works.

Generating Embeddings

Next, we need to generate our embeddings from our collected documents.

What are embeddings?

Embeddings are high-dimensional, continuous vector representations of text, words, or other types of data, where similar items have similar representations. They capture semantic relationships and features in a space where operations like distance or angle measurement can indicate similarity or dissimilarity.

In machine learning, embeddings are used to convert categorical, symbolic, or textual data into a form that algorithms can process more effectively, enabling tasks like natural language processing, recommendation systems, and more sophisticated pattern recognition.

With our textual data collected, we’ll be using OpenAI and Langchain to generate our embeddings. There are lots of different ways to generate embeddings (plenty packages that run locally, too), but using OpenAI API to get them is fast and easy for us. (and also dirt cheap)

NOTE: In a true production system, there would be much more consideration taken around generating embeddings. This is arguably the most important step in a RAG based system. We’d need to do experimentation with chunk size to see what gives us the best results. We’d need to explore our vectors to make sure their working as expected, remove noise, etc.

Generating embeddings using LangChain and OpenAI

Langchain makes it very easy to do create embeddings and store them in Redis without much thought, but this step requires extreme care to generate good results in a production system

The snippet below takes our scraped wikipedia sections, generates embeddings for them using OpenAI’s embeddings API, and stores them in Redis. Again, LangChain abstracts away a ton of complexity and makes this really easy for us.

 1from langchain.embeddings import OpenAIEmbeddings
 2
 3embeddings = OpenAIEmbeddings()
 4rds = Redis.from_texts(
 5    texts,
 6    embeddings,
 7    metadatas=metadata,
 8    redis_url="redis://localhost:6379",
 9    index_name="cheese",
10)

Implementation of RAG

Our RAG operates by creating an embedding of the user’s question and then finding the most semantically similar documents in our database (via cosine similarity between the embedding of our user’s query and the N closest documents in our database).

We then include these documents / snippets in our request to the LLM, telling it that they are the most relevant documents based on a similarity search. The LLM can then use these documents as reference when generating its response.

Here’s a simplified overview of the process with code snippets:

Embedding User Queries

The user’s query is converted into an embedding using the OpenAI API. This embedding represents the semantic content of the query in a format that can be compared against the pre-computed embeddings of the database articles.

1from langchain.embeddings import OpenAIEmbeddings
2
3embeddings = OpenAIEmbeddings()
4query_embedding = embeddings.embed_text(user_query)

Retrieving Related Articles

We then use the query embedding to perform a similarity search in the Redis database. It retrieves a set number of articles that are most semantically similar to the query.

1def get_related_articles(query_embedding, k=3):
2    return rds.similarity_search(query_embedding, k=k)

Integrating Retrieved Data into GPT-4 Prompts

The retrieved articles are formatted and integrated into the prompt for GPT-4. This allows GPT-4 to use the information from these articles to generate a response that is not only contextually relevant but also rich in content.

1def create_prompt_with_articles(query, articles):
2    article_summaries = [f"{article['title']}: {article['summary']}" for article in articles]
3    return f"Question: {query}\n\nRelated Information:\n" + "\n".join(article_summaries)

Generating the Response

Finally, the enriched prompt is fed to GPT-4, which generates a response based on both the user’s query and the additional context provided by the retrieved articles.

1response = openai.ChatCompletion.create(
2model="gpt-4",
3messages=[{"role": "system", "content": create_prompt_with_articles(user_query, related_articles)}]
4)

Through this process, CheeseGPT effectively combines the generative power of GPT-4 with the information retrieval capabilities of RAG, resulting in responses that are informative, accurate, and contextually rich.

The Chat Interface

CheeseGPT’s chat interface is an important component, orchestrating the interaction between the user, the retrieval-augmented generation system, and the underlying Large Language Model (LLM).

For the purposes of our example, we have built the bindings for the interface, but did not create a fully interactive interface.

Let’s dive into the key functions that make this interaction possible.

Connecting to Redis

1def rds_connect():
2    rds = Redis.from_existing_index(
3        embeddings,
4        redis_url="redis://localhost:6379",
5        index_name="cheese",
6        schema="redis_schema.yaml",
7    )
8    return rds

This function establishes a connection to the Redis database, where the precomputed embeddings of cheese-related Wikipedia pages are stored.

Applying Filters

1def get_filters():
2    is_not_external_link = RedisFilter.text("section") != 'External Links'
3    is_not_see_also = RedisFilter.text("section") != 'See Also'
4    _filter = is_not_external_link & is_not_see_also
5    return _filter

Filters are applied to ensure that irrelevant sections like ‘External Links’ and ‘See Also’ are excluded from the search results.

Deduplicating Results

1def dedupe_results(results):
2    seen = set()
3    deduped_results = []
4    for result in results:
5        if result.page_content not in seen:
6            deduped_results.append(result)
7            seen.add(result.page_content)
8    return deduped_results

This function ensures that duplicate content from the search results is removed, enhancing the quality of the final output. (This is necessary in our case because we were greedy / lazy when pulling our data / generating our vectors)

Retrieving Document Results

1def get_results(rds, question, k=3):
2    _filters = get_filters()
3    results = dedupe_results(rds.similarity_search(question, k=k, filter=_filters))
4    return results

This key function performs a similarity search in the Redis database using the user’s query, filtered and deduplicated.

Formatting RAG Results

1def format_rag_results(rag_results):
2    divider = "*********************RESULT*********************\n"
3    return [f"{divider} {result.page_content} ({result.metadata['url']}#{result.metadata['section']})" for result in
4            rag_results]

The function formats the search results, making them readable and including the source information for transparency.

This is what our message looks like when we send it to GPT-4. Our system prompt is first and includes instructions for the model to use the retrieved documents when answering the question.

In our user message, you can see the user’s question, and then the documents we retrieved, presented as a list with some formatting.

[
   {
      "role": "system",
      "content": "You are cheeseGPT, a retrieval augmented chatbot with expert knowledge of cheese. You are here to answer questions about cheese, and you should, when possible, cite your sources with the documents provided to you."
   },
   {
      "role": "user",
      "content": "User question: what is the biggest cheese sporting event.  Retrieved documents: ['*********************RESULT*********************\\n The Lucerne Cheese Festival (German: Käsefest Luzern) is a cheese festival held annually in Lucerne, Switzerland. It was established in 2001 and is normally run on a weekend in the middle of October at the Kapellplatz (Chapel Square) in the city centre. The next festival is planned to take place on 14 October 2023.The event features the biggest cheese market in central Switzerland, and offers the greatest selection of cheeses. As well as the cheese market and live demonstrations of cheesemaking, typical events during the festival include a milking competition and music such as the Swiss alphorn.The 2012 event featured over 200 varieties of cheese over 23 market stalls, including goat and sheep cheese. The 2020 event was almost cancelled because of social distancing restrictions during the COVID-19 pandemic, but was approved a few days before with a strict requirement to wear masks. Instead of the Kapellplatz, the festival was run from the nearby Kurplatz (Spa Square). 288 variety of cheeses were available at the festival, including cheesemakers from outside the local region such as the Bernese Jura and Ticino, who had their own festivals cancelled. Around 5,800 people attended the festival, lower than the previous year, with around two-thirds fewer sales. The following year\\'s event continued restrictions, where customers had to taste and buy cheese at a distance, though masks were no longer mandatory. The 2022 event featured demonstrations of the cheese making process, a chalet built of Swiss cheese, and a \"cheese chalet\" hosting cheese fondue and raclette.India Times in 2014 called it out as one of 10 world food festivals for foodies. (https://en.wikipedia.org/wiki/Lucerne_Cheese_Festival#Summary)', '*********************RESULT*********************\\n In addition to sampling and purchasing more than 4,600 cheeses in the Cheese Pavilion, visitors to the show are treated to various attractions throughout the day including cheese making demonstrations, trophy presentations and live cookery demonstrations. (https://en.wikipedia.org/wiki/International_Cheese_Awards#Show features)', '*********************RESULT*********************\\n The first annual event was held in 2000 in Oxfordshire, and was founded by Juliet Harbutt. Each year it is preceded by the British Cheese Awards, a ceremony which Harbutt created in 1994, judged by food experts and farmers, in which the best cheeses are awarded bronze, silver and gold medals.\\nAll cheeses are tasted blind, and the winners can then display their awards during the public-attended festival.  There are usually a variety of events at the festival such as seminars, masterclasses, and cheesemaking demonstrations.The event moved to Cheltenham in Gloucestershire in 2005. In 2006 the Sunday of the weekend was cancelled at great cost after the venue experienced flooding, and the decision was made to return to Oxfordshire in 2007.Early in 2008 the festival was sold to Cardiff Council; subsequently the event has been held in the grounds of Cardiff Castle in 2008, 2009, 2010, and 2011.The 2012 Great British Cheese Festival was held at Cardiff Castle on Saturday, September 22, and Sunday, September 23.\\nIn 2015, Harbutt returned to her native New Zealand. (https://en.wikipedia.org/wiki/The_Great_British_Cheese_Festival#History)']"
   }
]

Generating Messages for the LLM

1def get_messages(question, rag_results):
2    messages = [
3        {"role": "system", "content": system_prompt},
4        {"role": "user",
5         "content": f"User question: {question}.  Retrieved documents: {format_rag_results(rag_results)}"}
6    ]
7    return messages

This function prepares the input for the LLM, combining the system prompt, user question, and the retrieved documents.

The integration of these functions creates a seamless flow from the user’s question to the LLM’s informed response, enabling CheeseGPT to provide expert-level insights into the world of cheese.

Putting it all together might look something like…

1question = "what is the biggest cheese sporting event"
2results = get_results(rds, question, k=3)
3messages = get_messages(question, results)
4print(json.dumps(messages, indent=4))
5response = openai.ChatCompletion.create(
6    model="gpt-4",
7    messages=messages,
8)
9print(response)

Outcomes

So, let’s compare a question using our system vs. asking ChatGPT. We’ll use the same question above.

Using our system, we get this response:

 1The Lucerne Cheese Festival in Switzerland and the International Cheese Awards are two of the biggest cheese sporting events. 
 2
 3The Lucerne Cheese Festival, a yearly event in Lucerne, Switzerland, features the largest cheese market in central Switzerland 
 4and provides a vast selection of cheeses. The festival also includes live cheesemaking demonstrations, a milking competition, 
 5and music. In past years, it has presented over 200 varieties of cheese over 23 market stalls, including goat and sheep cheese. 
 6The festival attracted around 5,800 attendees in 2020, despite social distancing restrictions. The India Times recognized it 
 7as one of 10 world food festivals for foodies in 2014 [^1^].
 8
 9On the other hand, the International Cheese Awards allow visitors 
10to sample and purchase over 4,600 different cheeses in the Cheese Pavilion. The show features various attractions such as
11cheese making demonstrations, trophy presentations, and live cookery demonstrations [^2^].
12
13The Great British Cheese Festival, although not as large as the other two events, is also a notable cheese event. It 
14hosts the British Cheese Awards, where the best cheeses are awarded bronze, silver, and gold medals, judged by food 
15experts and farmers. All cheeses are tasted blindly, and the winners can display their awards during the festival. 
16The festival typically includes seminars, masterclasses, and cheesemaking demonstrations [^3^].
17
18In conclusion, both the Lucerne Cheese Festival and the International Cheese Awards are considered some of the biggest 
19cheese sporting events worldwide.
20
21[^1^]: (https://en.wikipedia.org/wiki/Lucerne_Cheese_Festival#Summary)
22[^2^]: (https://en.wikipedia.org/wiki/International_Cheese_Awards#Show_features)
23[^3^]: (https://en.wikipedia.org/wiki/The_Great_British_Cheese_Festival#History)"

And if we ask ChatGPT the same question…

I’m not sure which of these answer is more correct, and it doesn’t matter for the purposes of this example. The point is that we were able to retrieve and include our own information, external to the model, and make it use that information in it’s response.
It’s clear how many amazing use cases there are for something like this!

Hopefully this high level toy example was able to shed some light on what a RAG based system may look like. Checkout the additional resources linked before for more in-depth information. Thanks for reading!

Related reading

• RAG: From Context Injection to Knowledge Integration — why this pattern works, and the architectural limits it runs into.
• The Case Against pgvector — this toy used an in-memory vector store; here’s what changes when you put one in production.
• A Production Framework for LLM Feature Evaluation — deciding which LLM features are worth building at all.

Additional resources

https://github.com/ray-project/llm-applications/blob/main/notebooks/rag.ipynb https://github.com/pchunduri6/rag-demystified https://www.anyscale.com/blog/a-comprehensive-guide-for-building-rag-based-llm-applications-part-1

And yet, all too frequently, end users will go out of their way to completely avoid even attempting to read the error message.

And after a long hard day writing bug free code and solving customer problems, there is nothing less satisfying than a story that contains a screenshot of an error message that explains exactly the action the user needs to take the resolve the mistake they have made.

Below we present a simple and effective approach to handling errors in client-side fleshware.

An error message that requires interaction

If it’s not yet clear, I’ll take this opportunity to say that this is (mostly) a joke.

I’m certainly no UI designer, but with a little love, I expect this to be a core product feature in all modern softwares going forward.

Checkout the fiddle below to make your own improvements! https://jsfiddle.net/eg142dkp/2/

Wondering how I made this?

First I had to learn javascript. Then I had to learn CSS and HTML. After all that, I had to learn about Hugo and Shortcodes

By that point, I was a certified full-stack developer, and I was able to create the following shortcode to render the modal

 1<style>
 2    .modal {
 3        display: none; 
 4        position: fixed;
 5        z-index: 1; 
 6        left: 0;
 7        top: 0;
 8        width: 100%; 
 9        height: 100%; 
10        overflow: auto; 
11        background-color: rgb(0, 0, 0); 
12        background-color: rgba(0, 0, 0, 0.5); 
13
14    }
15
16    .modal-content {
17        display: block;
18        background-color: #fefefe;
19        margin: 25% auto; 
20        padding: 20px;
21        border: 3px solid #dc3545;
22        width: 60%; 
23        height: auto;
24        opacity: 1;
25        border-radius: 25px;
26        text-align: center;
27    }
28
29    .modal input {
30        background-color: #ffffff;
31        border: 1px solid #dc3545;
32        border-radius: 5px;
33        width: 50%;
34
35    }
36
37    .error-button {
38        background-color: #dc3545;
39        border: 1px solid #dc3545;
40        border-radius: 5px;
41        width: 40%;
42        color: #ffffff;
43        font-weight: bold;
44        display: grid;
45        margin: 0 auto;
46    }
47
48</style>
49<div id="errorModal" class="modal">
50
51    <div class="modal-content">
52        <div style="color: #dc3545;" id='errorMessage'></div>
53        <br>
54        <p style="color: gray; opacity: 65%">
55            type the error message to confirm you've seen it and disable the error modal
56        </p>
57        <input type="text" id="textinput" name="fname">
58    </div>
59
60</div>
61<div class="error-button">
62    <button type="submit" onclick="errorMessage('There is an error due to the way you did it')">
63        Click here to display error
64    </button>
65</div>
66<script>
67    function errorMessage(displayErrorMessage) {
68        let errorModal = document.getElementById("errorModal");
69        let errorMessage = document.getElementById('errorMessage')
70        errorMessage.textContent = displayErrorMessage
71        errorMessage.style.fontWeight = 'bold';
72        errorModal.style.display = "block";
73
74        let inputBox = document.getElementById('textinput');
75
76        inputBox.onkeyup = function () {
77            var text = inputBox.value
78            if (displayErrorMessage.startsWith(text)) {
79                errorMessage.innerHTML = "<span style='color: " + 'black' + ";'>" + text + "</span>" +
80                    "<span style='color: " + '#dc3545' + ";'>" +
81                    displayErrorMessage.replace(text, '') + "</span>"
82            }
83
84            if (text === displayErrorMessage) {
85                errorModal.style.display = "none";
86                inputBox.value = '';
87            }
88        }
89        console.log(errorMessage)
90    }
91</script>

and it goes right here in my layouts/shortcodes directory

1.
2└── layouts
3    └── shortcodes
4        └── error-modal.html

and then in my markdown file, I can just call it like this

1{{< error-model >}}

Hugo is neat

Update 12-27-22 – OpenAI

It is worth updating this post to say that I gave this prompt to the openai code-davinci-002 model

1/* create a button that says 'click to display error message'. 
2when the button is clicked, a modal should pop up that has an error message and a text box. 
3the modal should not have an x to close it. the only way to close the model is to type the error message into the text box */

and within seconds had a fully working implementation

Checkout openAI’s fiddle for comparison. Maybe we’re all wasting our time

https://jsfiddle.net/Lkyxuc0v/

 1var button = document.createElement('button');
 2button.innerHTML = 'click to display error message';
 3document.body.appendChild(button);
 4button.onclick = function () {
 5    var modal = document.createElement('div');
 6    modal.style.position = 'fixed';
 7    modal.style.top = '0';
 8    modal.style.left = '0';
 9    modal.style.width = '100%';
10    modal.style.height = '100%';
11    modal.style.backgroundColor = 'rgba(0,0,0,0.5)';
12    modal.style.zIndex = '9999';
13    document.body.appendChild(modal);
14    var errorMessage = document.createElement('div');
15    errorMessage.innerHTML = 'Error: You must type the error message into the text box to close this modal.';
16    errorMessage.style.position = 'absolute';
17    errorMessage.style.top = '50%';
18    errorMessage.style.left = '50%';
19    errorMessage.style.transform = 'translate(-50%, -50%)';
20    errorMessage.style.backgroundColor = 'white';
21    errorMessage.style.padding = '20px';
22    errorMessage.style.border = '1px solid black';
23    modal.appendChild(errorMessage);
24    var textBox = document.createElement('input');
25    textBox.type = 'text';
26    textBox.style.position = 'absolute';
27    textBox.style.top = '60%';
28    textBox.style.left = '50%';
29    textBox.style.transform = 'translate(-50%, -50%)';
30    textBox.style.width = '200px';
31    textBox.style.height = '20px';
32    textBox.style.border = '1px solid black';
33    modal.appendChild(textBox);
34    textBox.onkeyup = function () {
35        if (textBox.value === 'Error: You must type the error message into the text box to close this modal.') {
36            document.body.removeChild(modal);
37        }
38    }
39}

In the last post, we left off with a working reverse proxy, but we couldn’t access Jupyter lab due to its auth enforcement. Because of how we’re setting this up, we will be handling authentication upstream of Jupyter Lab, and we don’t want to rely on them for handling authentication. What we are going to do here is generally considered “unsafe.”
Again, if you’re looking to do this for your team, check out Jupyter Hub–it probably makes more sense for your use case.

To disable token auth, we will update our Jupyter Lab config.

There is an extensive config file for Jupyter Lab. In a production environment, I recommend using it (you can generate a sample file by running jupyter notebook --generate-config). But, for this toy example, we will pass our config as cmd line args. To disable token auth and to allow same-origin requests, we’re going to update our Jupyter Lab Dockerfile Entrypoint to include these arguments

"--ServerApp.token=", "--ServerApp.password=", "--ServerApp.allow_origin", "*"

Our Dockerfile should now look like

FROM ubuntu:20.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y \
      python3-pip \
      python3-dev && \
      python3 -m pip install jupyterlab

RUN useradd -ms /bin/bash jupyter
EXPOSE 8888
ENTRYPOINT ["jupyter", "lab", "--ip=0.0.0.0", "--port", "8888", "--allow-root", "--ServerApp.token=", "--ServerApp.password=", "--ServerApp.allow_origin", "*"]

And if we rebuild and start our docker compose again

docker compose build && docker compose up

We now get through to Juypter!

But if we try and open the Python kernel, we’ll notice it’s having trouble connecting.

Opening our browser dev tools shows that there is an issue with how our proxy is handling WebSockets

We’ll have to update our Nginx config to address this.
We will add these lines to set headers properly for WebSockets to our / location in the server block.

...
  server {
    listen       8000;
    server_name  localhost;

    location / {
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_hide_header "X-Frame-Options";
        proxy_pass http://upstream_jupyter;
        
        # websocket support
        proxy_http_version 1.1;
        proxy_set_header Upgrade "websocket";
        proxy_set_header Connection "Upgrade";
        proxy_read_timeout 86400;
    }
  }
...

And now, if we restart our containers using the updated config, we’ll see our kernel connects!

If you’re wondering how we will handle security when we’re basically giving whoever is using this a terminal into our cloud, the answer is using AWS to isolate the instance via IAM roles/ policy. We aren’t going to get too much into that in this post, but it is a valid concern. There isn’t much we can do to prevent a privilege escalation/container escape from a sophisticated user, but we can at least not give root access.

We’re going to update our Jupyter Dockerfile to have a new user, ‘jupyter’, and we’ll run Jupyter Lab as this user.

...
RUN useradd -ms /bin/bash jupyter
USER jupyter
...

We’re also going to update our ENTRYPOINT, so the Jupyter Lab root directory is set to the Jupyter user’s home directory

"--ServerApp.root_dir", "/home/jupyter", "--ServerApp.notebook_dir", "/home/jupyter"

Our Dockerfile should now look like

FROM ubuntu:20.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y \
      python3-pip \
      python3-dev && \
      python3 -m pip install jupyterlab

# add user and switch to them
RUN useradd -ms /bin/bash jupyter
USER jupyter

EXPOSE 8888
ENTRYPOINT ["jupyter", "lab", "--ip=0.0.0.0", "--port", "8888", "--ServerApp.token=", "--ServerApp.password=", "--ServerApp.allow_origin", "*", "--ServerApp.root_dir", "/home/jupyter", "--ServerApp.notebook_dir", "/home/jupyter"]

If we reload our site, we’ll see that the working directory is now set to /home/jupyter, and if we try to write to /, we’ll get a permissions error. It’s important to note that while this makes it a little more difficult for a malicious user to take over this ‘instance’, we will be giving them access to the internet, the ability to download and install packages, execute code, etc. It would not be too difficult for someone with mal intent to get around this. Changing the user and working directory does more to help an innocent user from accidentally breaking something.

Great! Now we have disabled token authentication, added a system user (who is now running Jupyter), and changed our notebook directory to our user’s directory! In the next post, we’ll set up a task definition and deploy to ECS.

Jupyter Lab is an open source web-based IDE for notebooks with Python and R support, geared towards the data science crowd. It’s a powerful, mature application with a potentially complex configuration. Our requirement was to deliver Jupyter Lab to users so that each user would have their own isolated “instance .” There is an off-the-shelf solution for this called Jupyter Hub that probably makes the most sense for your organization. This example will be a proof of concept on how you could roll your solution.

Step 1–Jupyter Lab Docker

If you aren’t familiar with Docker, we’re going to be using it a lot here, so check out some guides

Our first step will be getting Jupyter Lab up and running in a container. There are many Docker images available on (Docker hub)[https://hub.docker.com/] for Jupyter Lab, but since we’re rolling everything ourselves, we might as well make our own image. It also gives us more control over our code–it’s also a pretty simple Dockerfile.

FROM ubuntu:20.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y \
      python3-pip \
      python3-dev \
    && python3 -m pip install jupyterlab

EXPOSE 8888
ENTRYPOINT ["jupyter", "lab", "--ip=0.0.0.0", "--port", "8888", "--allow-root"]