xbe.at

LLMs and generative AI / RAG and retrieval

Deploying Private RAG APIs with LitServe

11 min read, 15 sections

Contents
  1. Introduction to LitServe and RAG
  2. Understanding RAG (Retrieval-Augmented Generation)
  3. Configuring vLLM for Efficient LLM Serving
  4. Setting Up Qdrant Vector Database
  5. Implementing a Custom RAG Pipeline
  6. Securing Your Private API
  7. Optimizing RAG Performance
  8. Handling Multi-Modal Inputs
  9. Implementing Continuous Learning
  10. Real-Life Example: Customer Support Bot
  11. Real-Life Example: Personalized Learning Assistant
  12. Handling Edge Cases and Error Scenarios
  13. Monitoring and Analytics for RAG Systems
  14. Scaling RAG Systems for High-Load Scenarios
  15. Additional Resources

Introduction to LitServe and RAG

LitServe is a scalable, high-performance inference engine for AI models with a minimal interface. It enables the deployment of Retrieval-Augmented Generation (RAG) applications behind a private API. This setup is framework-agnostic and fully customizable, allowing for efficient LLM serving and high-performance vector search.

# Example of a basic LitServe setup
from litserve import LitServe
from vllm import LLMEngine
from qdrant_client import QdrantClient

# Initialize LitServe
litserve = LitServe()

# Set up LLM engine (using vLLM for efficiency)
llm_engine = LLMEngine.from_pretrained("llama-3.2")

# Initialize Qdrant vector database
qdrant_client = QdrantClient(host="localhost", port=6333)

# Register components with LitServe
litserve.register_llm(llm_engine)
litserve.register_vectordb(qdrant_client)

# Start the server
litserve.serve()

Understanding RAG (Retrieval-Augmented Generation)

RAG is a technique that enhances language models by retrieving relevant information from a knowledge base before generating responses. This approach combines the strengths of retrieval-based and generative models, resulting in more accurate and contextually relevant outputs.

# Simplified RAG process
def rag_process(query, knowledge_base, llm):
    # Retrieve relevant documents
    relevant_docs = knowledge_base.search(query)
    
    # Augment the query with retrieved information
    augmented_query = f"{query}\nContext: {' '.join(relevant_docs)}"
    
    # Generate response using the augmented query
    response = llm.generate(augmented_query)
    
    return response

# Usage
result = rag_process("What is the capital of France?", knowledge_base, llm)
print(result)
# Output: The capital of France is Paris.

Configuring vLLM for Efficient LLM Serving

vLLM is a library for efficient LLM inference, offering optimizations like continuous batching and PagedAttention. It significantly improves throughput and reduces latency compared to traditional serving methods.

from vllm import LLMEngine, SamplingParams

# Initialize the LLM engine with Llama 3.2
engine = LLMEngine.from_pretrained("llama-3.2")

# Set up sampling parameters
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=100
)

# Generate text
prompt = "Explain the concept of artificial intelligence:"
outputs = engine.generate([prompt], sampling_params)

# Process the output
for output in outputs:
    generated_text = output.outputs[0].text
    print(generated_text)

Setting Up Qdrant Vector Database

Qdrant is a high-performance vector database optimized for similarity search and retrieval. It allows efficient storage and querying of high-dimensional vectors, making it ideal for RAG applications.

from qdrant_client import QdrantClient
from qdrant_client.http import models

# Initialize Qdrant client
client = QdrantClient("localhost", port=6333)

# Create a collection for storing document embeddings
client.create_collection(
    collection_name="documents",
    vectors_config=models.VectorParams(size=768, distance=models.Distance.COSINE)
)

# Add vectors to the collection
client.upsert(
    collection_name="documents",
    points=[
        models.PointStruct(
            id=1, 
            vector=[0.05, 0.61, 0.76, 0.74],
            payload={"text": "Sample document 1"}
        ),
        models.PointStruct(
            id=2, 
            vector=[0.19, 0.81, 0.75, 0.11],
            payload={"text": "Sample document 2"}
        )
    ]
)

# Perform a similarity search
search_result = client.search(
    collection_name="documents",
    query_vector=[0.2, 0.1, 0.9, 0.7],
    limit=1
)

print(search_result[0].payload["text"])
# Output: Sample document 1

Implementing a Custom RAG Pipeline

Creating a custom RAG pipeline involves combining the LLM and vector database to retrieve relevant information and generate augmented responses.

from vllm import LLMEngine, SamplingParams
from qdrant_client import QdrantClient

class RAGPipeline:
    def __init__(self, llm_model, vector_db):
        self.llm = LLMEngine.from_pretrained(llm_model)
        self.vector_db = QdrantClient(vector_db)
    
    def retrieve(self, query, k=3):
        # Assume we have a function to convert query to vector
        query_vector = query_to_vector(query)
        results = self.vector_db.search(
            collection_name="documents",
            query_vector=query_vector,
            limit=k
        )
        return [r.payload["text"] for r in results]
    
    def generate(self, query, context):
        prompt = f"Query: {query}\nContext: {' '.join(context)}\nAnswer:"
        sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
        output = self.llm.generate([prompt], sampling_params)
        return output[0].outputs[0].text
    
    def rag_query(self, query):
        context = self.retrieve(query)
        return self.generate(query, context)

# Usage
rag = RAGPipeline("llama-3.2", "localhost:6333")
response = rag.rag_query("What is the capital of France?")
print(response)
# Output: Based on the context provided, the capital of France is Paris.

Securing Your Private API

When deploying a private RAG API, it's crucial to implement proper security measures to protect your data and resources.

from flask import Flask, request, jsonify
from functools import wraps

app = Flask(__name__)

# Simulated user database
users = {
    "alice": "password123",
    "bob": "securepass456"
}

def authenticate(username, password):
    return username in users and users[username] == password

def require_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not authenticate(auth.username, auth.password):
            return jsonify({"message": "Authentication required"}), 401
        return f(*args, **kwargs)
    return decorated

@app.route('/rag_query', methods=['POST'])
@require_auth
def rag_query():
    data = request.json
    query = data.get('query')
    if not query:
        return jsonify({"error": "Query is required"}), 400
    
    # Process the RAG query here
    result = "Simulated RAG response"
    
    return jsonify({"result": result})

if __name__ == '__main__':
    app.run(ssl_context='adhoc')  # Enable HTTPS

Optimizing RAG Performance

To improve the performance of your RAG system, consider techniques such as caching, query preprocessing, and efficient vector indexing.

import time
from functools import lru_cache

class OptimizedRAGPipeline:
    def __init__(self, llm_model, vector_db):
        self.llm = LLMEngine.from_pretrained(llm_model)
        self.vector_db = QdrantClient(vector_db)
    
    @lru_cache(maxsize=1000)
    def cached_retrieve(self, query):
        return self.retrieve(query)
    
    def retrieve(self, query, k=3):
        # Implementation as before
        pass
    
    def preprocess_query(self, query):
        # Simple preprocessing: lowercase and remove punctuation
        return ''.join(c.lower() for c in query if c.isalnum() or c.isspace())
    
    def rag_query(self, query):
        start_time = time.time()
        
        preprocessed_query = self.preprocess_query(query)
        context = self.cached_retrieve(preprocessed_query)
        response = self.generate(preprocessed_query, context)
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        return response, processing_time

# Usage
rag = OptimizedRAGPipeline("llama-3.2", "localhost:6333")
response, time_taken = rag.rag_query("What is the capital of France?")
print(f"Response: {response}")
print(f"Time taken: {time_taken:.2f} seconds")

Handling Multi-Modal Inputs

Extend your RAG system to handle various input types, such as text, images, or audio, to create a more versatile and powerful application.

import base64
from PIL import Image
import io

class MultiModalRAG:
    def __init__(self, llm_model, vector_db):
        self.llm = LLMEngine.from_pretrained(llm_model)
        self.vector_db = QdrantClient(vector_db)
    
    def process_text(self, text):
        # Existing text processing logic
        pass
    
    def process_image(self, image_data):
        # Convert base64 to image
        image = Image.open(io.BytesIO(base64.b64decode(image_data)))
        # Perform image analysis (e.g., object detection, OCR)
        # Return textual description or extracted text
        return "Image analysis result"
    
    def process_audio(self, audio_data):
        # Convert audio to text using speech recognition
        # This is a placeholder for actual speech-to-text logic
        return "Transcribed audio text"
    
    def rag_query(self, query, input_type='text'):
        if input_type == 'text':
            context = self.process_text(query)
        elif input_type == 'image':
            context = self.process_image(query)
        elif input_type == 'audio':
            context = self.process_audio(query)
        else:
            raise ValueError("Unsupported input type")
        
        return self.generate(context)

# Usage
rag = MultiModalRAG("llama-3.2", "localhost:6333")
text_response = rag.rag_query("What is the capital of France?", input_type='text')
image_response = rag.rag_query(base64_encoded_image, input_type='image')
audio_response = rag.rag_query(base64_encoded_audio, input_type='audio')

Implementing Continuous Learning

Enhance your RAG system with continuous learning capabilities to improve its knowledge base over time.

import datetime

class ContinuousLearningRAG:
    def __init__(self, llm_model, vector_db):
        self.llm = LLMEngine.from_pretrained(llm_model)
        self.vector_db = QdrantClient(vector_db)
        self.feedback_log = []
    
    def rag_query(self, query):
        # Existing RAG query logic
        pass
    
    def log_feedback(self, query, response, feedback):
        self.feedback_log.append({
            'query': query,
            'response': response,
            'feedback': feedback,
            'timestamp': datetime.datetime.now()
        })
    
    def update_knowledge_base(self):
        for entry in self.feedback_log:
            if entry['feedback'] == 'positive':
                # Add the query-response pair to the knowledge base
                self.add_to_knowledge_base(entry['query'], entry['response'])
        
        # Clear the feedback log after processing
        self.feedback_log.clear()
    
    def add_to_knowledge_base(self, query, response):
        # Convert query-response pair to vector and add to vector database
        vector = text_to_vector(f"{query} {response}")
        self.vector_db.upsert(
            collection_name="documents",
            points=[
                models.PointStruct(
                    id=generate_unique_id(),
                    vector=vector,
                    payload={"text": f"Q: {query} A: {response}"}
                )
            ]
        )

# Usage
rag = ContinuousLearningRAG("llama-3.2", "localhost:6333")
response = rag.rag_query("What is the capital of France?")
rag.log_feedback("What is the capital of France?", response, "positive")
rag.update_knowledge_base()

Real-Life Example: Customer Support Bot

Implement a customer support bot using the RAG system to provide accurate and context-aware responses to customer queries.

class CustomerSupportRAG:
    def __init__(self, llm_model, vector_db):
        self.rag = RAGPipeline(llm_model, vector_db)
        self.product_kb = self.load_product_knowledge()
    
    def load_product_knowledge(self):
        # Load product information from a database or file
        return {
            "ProductA": {"features": ["Feature1", "Feature2"], "price": 99.99},
            "ProductB": {"features": ["Feature3", "Feature4"], "price": 149.99}
        }
    
    def handle_query(self, query):
        # Check if the query is about a specific product
        for product, info in self.product_kb.items():
            if product.lower() in query.lower():
                context = f"Product: {product}\nFeatures: {', '.join(info['features'])}\nPrice: ${info['price']}"
                return self.rag.generate(query, [context])
        
        # If not about a specific product, use general RAG
        return self.rag.rag_query(query)

# Usage
support_bot = CustomerSupportRAG("llama-3.2", "localhost:6333")
response = support_bot.handle_query("What are the features of ProductA?")
print(response)
# Output: ProductA features include Feature1 and Feature2. It is priced at $99.99.

response = support_bot.handle_query("How do I return a product?")
print(response)
# Output: [General return policy information from the RAG system]

Real-Life Example: Personalized Learning Assistant

Create a personalized learning assistant that uses RAG to provide tailored educational content and answer students' questions.

class LearningAssistantRAG:
    def __init__(self, llm_model, vector_db):
        self.rag = RAGPipeline(llm_model, vector_db)
        self.student_profiles = {}
    
    def update_student_profile(self, student_id, topic, performance):
        if student_id not in self.student_profiles:
            self.student_profiles[student_id] = {}
        self.student_profiles[student_id][topic] = performance
    
    def get_personalized_content(self, student_id, query):
        profile = self.student_profiles.get(student_id, {})
        context = f"Student profile: {profile}"
        
        topic = self.identify_topic(query)
        level = profile.get(topic, "beginner")
        
        augmented_query = f"Provide a {level}-level explanation for: {query}"
        response = self.rag.generate(augmented_query, [context])
        
        return response
    
    def identify_topic(self, query):
        topics = {
            "math": ["equation", "algebra", "geometry"],
            "science": ["physics", "chemistry", "biology"],
            "history": ["ancient", "modern", "world war"]
        }
        
        for topic, keywords in topics.items():
            if any(keyword in query.lower() for keyword in keywords):
                return topic
        return "general"

# Usage
assistant = LearningAssistantRAG("llama-3.2", "localhost:6333")
assistant.update_student_profile("student123", "math", "intermediate")
response = assistant.get_personalized_content("student123", "Explain quadratic equations")
print(response)

Handling Edge Cases and Error Scenarios

Implement robust error handling and edge case management in your RAG system to ensure reliability and graceful degradation.

class RobustRAGSystem:
    def __init__(self, llm_model, vector_db):
        self.rag = RAGPipeline(llm_model, vector_db)
        self.fallback_responses = {
            "no_context": "I'm sorry, I couldn't find relevant information to answer your query.",
            "llm_error": "I'm having trouble generating a response. Please try again later.",
            "invalid_query": "I didn't understand your query. Could you please rephrase it?"
        }
    
    def safe_query(self, query):
        try:
            if not self.is_valid_query(query):
                return self.fallback_responses["invalid_query"]
            
            context = self.safe_retrieve(query)
            if not context:
                return self.fallback_responses["no_context"]
            
            response = self.safe_generate(query, context)
            return response if response else self.fallback_responses["llm_error"]
        
        except Exception as e:
            logging.error(f"Error processing query: {str(e)}")
            return "An unexpected error occurred. Please try again later."
    
    def is_valid_query(self, query):
        return len(query.strip()) > 0 and len(query) < 1000
    
    def safe_retrieve(self, query):
        try:
            return self.rag.retrieve(query)
        except Exception as e:
            logging.error(f"Retrieval error: {str(e)}")
            return None
    
    def safe_generate(self, query, context):
        try:
            return self.rag.generate(query, context)
        except Exception as e:
            logging.error(f"Generation error: {str(e)}")
            return None

# Usage
robust_rag = RobustRAGSystem("llama-3.2", "localhost:6333")
response = robust_rag.safe_query("What is the capital of France?")
print(response)

Monitoring and Analytics for RAG Systems

Implement monitoring and analytics to track the performance and usage of your RAG system, enabling data-driven improvements.

import time
from collections import defaultdict

class MonitoredRAGSystem:
    def __init__(self, llm_model, vector_db):
        self.rag = RAGPipeline(llm_model, vector_db)
        self.query_count = 0
        self.total_latency = 0
        self.topic_distribution = defaultdict(int)
        self.error_counts = defaultdict(int)
    
    def monitored_query(self, query):
        start_time = time.time()
        self.query_count += 1
        
        try:
            topic = self.identify_topic(query)
            self.topic_distribution[topic] += 1
            
            response = self.rag.rag_query(query)
            
            end_time = time.time()
            self.total_latency += (end_time - start_time)
            
            return response
        
        except Exception as e:
            error_type = type(e).__name__
            self.error_counts[error_type] += 1
            raise
    
    def get_analytics(self):
        avg_latency = self.total_latency / self.query_count if self.query_count > 0 else 0
        return {
            "total_queries": self.query_count,
            "average_latency": avg_latency,
            "topic_distribution": dict(self.topic_distribution),
            "error_counts": dict(self.error_counts)
        }
    
    def identify_topic(self, query):
        # Simple topic identification logic
        topics = ["general", "science", "history", "technology"]
        return topics[hash(query) % len(topics)]

# Usage
monitored_rag = MonitoredRAGSystem("llama-3.2", "localhost:6333")

for _ in range(100):
    try:
        monitored_rag.monitored_query("Random query " + str(_))
    except:
        pass

print(monitored_rag.get_analytics())

Scaling RAG Systems for High-Load Scenarios

Design strategies for scaling your RAG system to handle high-load scenarios and maintain performance under increased query volumes.

import multiprocessing
from concurrent.futures import ThreadPoolExecutor

class ScalableRAGSystem:
    def __init__(self, llm_model, vector_db, num_workers=4):
        self.rag = RAGPipeline(llm_model, vector_db)
        self.num_workers = num_workers
        self.executor = ThreadPoolExecutor(max_workers=num_workers)
    
    def process_query_batch(self, queries):
        return list(self.executor.map(self.rag.rag_query, queries))
    
    def start_query_server(self, input_queue, output_queue):
        while True:
            batch = []
            while len(batch) < 10:  # Collect up to 10 queries
                try:
                    query = input_queue.get(timeout=0.1)
                    batch.append(query)
                except multiprocessing.queues.Empty:
                    break
            
            if batch:
                results = self.process_query_batch(batch)
                for result in results:
                    output_queue.put(result)
    
    def scale_out(self):
        input_queue = multiprocessing.Queue()
        output_queue = multiprocessing.Queue()
        
        processes = []
        for _ in range(self.num_workers):
            p = multiprocessing.Process(target=self.start_query_server, 
                                        args=(input_queue, output_queue))
            p.start()
            processes.append(p)
        
        return input_queue, output_queue, processes

# Usage
scalable_rag = ScalableRAGSystem("llama-3.2", "localhost:6333", num_workers=4)
input_q, output_q, workers = scalable_rag.scale_out()

# Simulate incoming queries
for i in range(100):
    input_q.put(f"Query {i}")

# Collect results
results = []
for _ in range(100):
    results.append(output_q.get())

print(f"Processed {len(results)} queries")

# Clean up
for p in workers:
    p.terminate()

Additional Resources

For those interested in diving deeper into the topics covered in this presentation, here are some valuable resources:

  1. "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks" by Lewis et al. (2020) ArXiv: https://arxiv.org/abs/2005.11401
  2. "Language Models are Few-Shot Learners" by Brown et al. (2020) ArXiv: https://arxiv.org/abs/2005.14165
  3. "Efficient Transformers: A Survey" by Tay et al. (2020) ArXiv: https://arxiv.org/abs/2009.06732
  4. "What Language Model to Train if You Have One Million GPU Hours?" by Scao et al. (2022) ArXiv: https://arxiv.org/abs/2206.15342

These papers provide in-depth insights into language models, retrieval-augmented generation, and efficient training techniques, which are fundamental to building advanced RAG systems.

More in RAG and retrieval

  1. EfficientRAG Supercharging Multi-Hop QA for LLMs12 min

    EfficientRAG is an advanced technique for enhancing multi-hop question answering (QA) in Large Language Models (LLMs) using Python. This method combines...

  2. Comparing RAG and Fine-Tuned Language Models15 min

    The fundamental architectural differences between RAG and fine-tuned models lie in how they process and generate responses. RAG systems employ a retrieval...

  3. End-to-End Retrieval-Augmented Generation (RAG) Application Using Haystack, Pinecone, and Mistral AI7 min

    End-to-End Retrieval-Augmented Generation (RAG) is an advanced natural language processing technique that combines the power of large language models with...

  4. Chunking Strategies for Improved RAG Performance19 min

    Text chunking is a critical preprocessing step in RAG systems that involves breaking down large documents into smaller, meaningful segments. This process...

  5. Enhancing LLM Context with Recursive Summarization Using Python6 min

    Large Language Models (LLMs) have limited context windows. Recursive summarization is a technique to extend this context by iteratively condensing...

  6. Auto-Document Retrieval for Improved RAG Systems10 min

    Auto-Document Retrieval is an advanced technique that enhances Retrieval-Augmented Generation (RAG) systems. It autonomously selects and retrieves relevant...