Building a Production-Grade Email AI Assistant Using RAG and Multi-Stage Retrieval

shortwave · Tech · 2023 · source Agentic: yes · promotional bias: med · source-backed · extracted with prompt 0.5.4 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: Production email assistant over personal inbox history

Shortwave wanted an AI executive assistant inside email that could answer personalized questions, compose in a user's voice, cite sources, and search an entire inbox despite heterogeneous data, ambiguous conversational context, prompt limits, privacy constraints, and low-latency product expectations.

Constraints:

• users' inboxes contain large, private, heterogeneous text data including correspondence, calendar invites, receipts, notifications, and newsletters
• many questions require personalized data retrieval while general-knowledge questions require no tools
• retrieved email history can exceed prompt limits by orders of magnitude
• email search must combine semantic meaning with dates, labels, contacts, recency, and other metadata
• the assistant needed source citations and clean UI post-processing
• the end-to-end system had to answer almost all questions within 3-5 seconds despite multiple LLM calls and ML/search infrastructure

Architecture

Shortwave built a four-step assistant around GPT-4 and modular tools. A tool-selection LLM chooses zero, one, or multiple tools based on the query and app context. Tools retrieve data in parallel, with the EmailHistory/AI Search tool running query reformulation, feature extraction, full-text/metadata search, embedding search through Pinecone, heuristic re-ranking, and MS Marco MiniLM cross-encoder re-ranking on Shortwave GPUs. The final answer is generated by a single GPT-4 call with selected context, followed by rich-text conversion, citations, and suggested actions.

Diagram confidence: explicit (topology described by the source).

flowchart LR
  shortwave-app[Shortwave app] --> tool-selection-llm[Tool selection LLM]
  tool-selection-llm --> tool-registry[Tool registry]
  tool-registry --> ai-search-tool[AI Search tool]
  ai-search-tool --> query-reformulator[Query reformulator]
  query-reformulator --> feature-extractors[Feature extractors]
  feature-extractors --> traditional-search-infrastructure[Traditional search infrastructure]
  feature-extractors --> pinecone-vector-database[Pinecone vector database]
  traditional-search-infrastructure --> heuristic-re-ranker[Heuristic re-ranker]
  pinecone-vector-database --> heuristic-re-ranker
  heuristic-re-ranker --> cross-encoder-re-ranker[Cross-encoder re-ranker]
  tool-registry --> compose-tool[Compose tool]
  tool-registry --> calendar-and-thread-tools[Calendar and thread tools]
  tool-registry --> final-answer-llm[Final answer LLM]
  final-answer-llm --> post-processing-layer[Post-processing layer]
  post-processing-layer --> shortwave-app

Components

Data flow

A user asks Shortwave AI a question. GPT-4 selects the necessary data tools from app state, conversation history, current thread, drafts, calendar, and email history. Selected tools run in parallel; EmailHistory reformulates the query, extracts search features, combines traditional search with Pinecone vector search, narrows and re-ranks fragments, and returns ordered context. Shortwave then sends the original question plus tool context to GPT-4 for one final answer, post-processes the output, adds citations, and suggests actions.

Why this architecture

Shortwave chose a modular tool architecture because email questions require different data sources, but kept the final reasoning in one large GPT-4 call because long LLM chains lost information and accumulated errors. The heavy retrieval complexity lives inside AI Search, which uses hybrid retrieval and two-stage re-ranking to squeeze a user's email history down to a few useful fragments fast enough for a product interaction.

Models used: GPT-4 · GPT-3 · Instructor · MS Marco MiniLM · fast LLM feature extractors

Scale

• The assistant answers almost all questions within 3-5 seconds.
  “Despite the complexity of the system, we have optimized it to be able to answer almost all questions within 3-5 seconds. This end-to-end latency includes multiple LLM calls, vector DB lookups, ElasticSearch queries, calls to open source ML models, and more.”article, stated
• AI Search may retrieve a thousand or more candidate emails before ranking.
  “Often, this results in a thousand or more results – far too many to fit into a prompt. We need to whittle this down to just a few dozen results, and we need to do this fast.”article, stated
• The system uses multiple LLM calls, vector DB lookups, ElasticSearch, open-source ML models, and GPU clusters within one request.
  “This end-to-end latency includes multiple LLM calls, vector DB lookups, ElasticSearch queries, calls to open source ML models, and more.”article, stated
• Feature extraction runs many independent fast-LLM calls in parallel.
  “We do this via a large number of parallel calls to a fast LLM, each optimized to extract a specific feature. Having these LLM calls be independent of each other enables a degree of modularity that makes it easier to test and evolve them.”article, stated
• Shortwave runs embedding and cross-encoder models on its own GPU-backed infrastructure.
  “We do this on our own servers with GPU acceleration for improved performance, security, and cost. Then, we store that embedding in our Vector database (Pinecone), namespaced per-user (the namespacing is one of the main reasons we chose Pinecone).”article, stated

  “We use MS Marco MiniLM, running on our own GPUs.”article, stated

Alternatives considered (roads not taken)

Break reasoning into long chains of LLM calls

Rejected because: Shortwave found long chains introduced data loss and errors, so the final answer comes from one GPT-4 call with all selected context.

Include all app data in every LLM call

Rejected because: Real-world context windows and loading latency required careful selection of what data to include.

Use traditional search infrastructure alone

Rejected because: Earlier attempts to turn email history into an actionable knowledge base with traditional search were underwhelming.

Trust semantic similarity score alone

Rejected because: Semantic scores miss metadata concerns like labels, dates, contacts, recency, and embedding-model inconsistencies.

Run the expensive cross-encoder over every retrieved email

Rejected because: The cross-encoder is smarter but slow, so Shortwave first narrows candidates with fast heuristics.

Patterns

tool-use-orchestration (assistant-architecture)

GPT-4 selects which modular tools are needed for a user question, allowing zero, one, or multiple sources such as current thread, email history, calendar, compose, or drafts.

rag-retrieval (email-search)

AI Search retrieves and ranks relevant email fragments from the user's full history so GPT-4 can answer with grounded inbox context.

query-rewriting (retrieval-preprocessing)

Context-dependent user questions are rewritten into standalone queries before retrieval.

hybrid-search (retrieval)

Shortwave combines full-text/metadata search with embedding-based semantic search and extracted features.

retrieval-ranking (retrieval)

Retrieved emails are narrowed through heuristic scoring and then cross-encoder relevance scoring before being included in the final prompt.

latency-optimization (production-performance)

The system uses parallel tool retrieval, parallel feature extraction, small-batch loading, streaming, pipelining, and GPU compute to hit 3-5 second latency.

Tradeoffs

Centralize final reasoning in one large GPT-4 call

Chose single final-answer prompt with selected context over multi-stage LLM reasoning chains.

Rationale: The team trusted GPT-4's reasoning and context window more than brittle chained summaries that could lose data.

Use modular tools for data sourcing

Chose calendar, email history, draft, compose, and current-thread tools over one monolithic assistant data loader.

Rationale: Tool boundaries made heterogeneous data sources easier to reason about, test, and evolve independently.

Run embeddings on Shortwave GPU servers and store them in Pinecone

Chose self-hosted Instructor embeddings plus per-user Pinecone namespaces over outsourcing all embedding computation or using a non-namespaced vector store.

Rationale: Self-hosting improved performance, security, and cost; Pinecone namespacing matched the privacy boundary of user email histories.

Use heuristic filtering before ML re-ranking

Chose fast local heuristics followed by cross-encoder scoring over cross-encoding every candidate email.

Rationale: The cross-encoder was more accurate but slow, so cheaper metadata-aware heuristics reduced the set first.

Allocate prompt tokens heuristically across tools

Chose heuristics when retrieved content exceeds prompt limits over trying to include all tool outputs equally.

Rationale: EmailHistory usually produces more content than can fit, so the system needs explicit tradeoffs among tools.

Gotchas (what bit them)

Long LLM chains lost information and compounded errors.

Root cause: Breaking reasoning into many LLM stages introduced data loss and mistakes at each step.

Resolution: Shortwave used one GPT-4 final-answer call with all selected context in a large prompt.

Personalized email questions often lack standalone search meaning.

Root cause: Users ask follow-up questions that depend on chat history, visible thread, or current draft context.

Resolution: The AI Search tool reformulates queries using screen and conversation context before retrieval.

Email retrieval can return too many candidates for a prompt.

Root cause: Traditional and embedding searches may produce a thousand or more relevant-looking results.

Resolution: Shortwave uses fast heuristic ranking, small-batch loading, length-capped chunks, and cross-encoder re-ranking to reduce candidates to useful fragments.

Semantic similarity alone misses metadata-intent alignment.

Root cause: Embedding scores represent content similarity but not labels, dates, contacts, recency, or feature confidence.

Resolution: Shortwave re-applies local heuristics before and after cross-encoder scoring.

The smartest re-ranker was too slow to use everywhere.

Root cause: MS Marco MiniLM cross-encoder scoring is more accurate than heuristics but computationally expensive.

Resolution: Only top-ranked fragments from the heuristic phase are passed through the cross-encoder.

Evaluation

The article describes production validation through alpha-user concept testing, correctness of summaries, tool-selection examples, confidence-scored feature extraction, retrieval/re-ranking stages, source citation post-processing, and end-to-end latency measurement.

Metrics: almost all questions answered within 3-5 seconds · confidence scores for extracted features · semantic similarity scores · cross-encoder relevance scores · source citations in final output · retrieval narrowed from 1000+ results to a few dozen useful fragments

Outcomes

• Shortwave launched a production RAG-based email assistant that could answer questions over a user's inbox.
• The assistant supports modular data tools, write-in-your-voice composition, source citations, and rich UI actions.
• AI Search combines query reformulation, feature extraction, full-text search, vector search, heuristic ranking, and cross-encoder ranking.
• The production system answers almost all questions within 3-5 seconds despite multiple LLM calls, vector DB lookups, ElasticSearch queries, open-source ML models, and GPU compute.
• The team framed the system as v1 and planned rapid iteration as models and context windows improved.

The system operates over entire user inbox histories, often narrows 1000+ retrieved results to prompt-sized fragments, and hits 3-5 second latency through concurrency, streaming, pipelining, beefy servers, and high-end GPU clusters.