Uber's Genie: A RAG-based AI Assistant for Engineering Support
This article details the system architecture of Uber's internal AI assistant, Genie, which leverages Retrieval-Augmented Generation (RAG) to answer engineering support questions. It focuses on how Genie integrates knowledge preparation, information retrieval using embeddings, and LLM-based answer generation to address the challenge of scattered documentation and reduce on-call burden in a large organization. The design highlights the ETL pipeline for knowledge ingestion, vector database for semantic search, and strategies to mitigate LLM hallucinations.
Read original on ByteByteGoThe Challenge of Internal Knowledge Management
Large organizations like Uber face significant challenges in managing and disseminating internal engineering knowledge. With numerous platform teams supporting thousands of developers, questions about configuration, errors, and features are constant. This often leads to platform teams spending considerable time on repetitive support via channels like Slack, rather than focusing on development. Existing documentation is frequently scattered across wikis and internal forums, making it difficult to find, thus encouraging direct human queries.
Genie's Architecture: Retrieval-Augmented Generation (RAG)
Uber's Genie addresses this by combining information retrieval with large language models (LLMs) using a Retrieval-Augmented Generation (RAG) approach. Unlike fine-tuning an LLM, RAG allows Genie to query an external knowledge base for relevant information at query time and then feed that information to a general-purpose LLM to generate an answer. This avoids the significant compute and data labeling overhead of fine-tuning and allows for easier knowledge updates.
Why RAG over Fine-Tuning?
RAG enables systems to utilize current, domain-specific knowledge without retraining the entire language model. It's particularly effective for rapidly changing data or proprietary information, as updates to the external knowledge base immediately improve retrieval without costly model updates. This significantly reduces time-to-market and operational costs compared to frequent fine-tuning.
Core Components and Workflow
- Knowledge Ingestion Pipeline (ETL): Uses Apache Spark to extract content from sources like Engwiki and internal Stack Overflow, transforms it by chunking documents and generating embeddings, and loads it into the search infrastructure. Crucially, source URLs are preserved for attribution.
- Embedding Generation: LangChain is used for text chunking, and an OpenAI embedding model generates numerical vector representations (embeddings) for each chunk. Related texts have similar embeddings, enabling semantic search even with different phrasing.
- Search Infrastructure: Uber's internal vector database, Sia, stores the embeddings, while Terrablob (blob storage) holds index data. A bootstrap job ingests data into Sia, and Spark jobs build and merge an efficient search index. Serving machines (leaves) download this index to handle live search requests.
- Knowledge Service: The central backend component. When a user asks a question in Slack, Genie sends it to the Knowledge Service. It converts the question into an embedding, queries Sia for relevant document chunks, and then sends these chunks along with the original question to an LLM to generate the final response.
Ensuring Accuracy and Feedback
To combat LLM hallucination, Genie explicitly provides the LLM with retrieved results as separate sub-contexts, each linked to its source URL. This provides boundaries for the model and allows users to verify information. If Genie cannot answer, it provides an option to ask a human. Uber also collects feedback and continuously evaluates Genie's performance by checking both retrieval quality (suitable passages) and generation quality (adequate answers).