scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

feat(cookbook): add multi-provider RAG gateway example with LlamaIndex, Haystack, and RAGAS

Browse source 
- Demonstrates LiteLLM as unified gateway for 100+ LLM providers
- Includes LlamaIndex RAG pipeline with semantic search
- Includes Haystack RAG pipeline with custom LiteLLM component
- Adds RAGAS evaluation for comparing models/providers
- Shows production patterns: retries, fallbacks, timeouts, logging
- Provides sample AI/ML knowledge base and evaluation dataset
- Comprehensive README with setup, usage, and customization guide
This commit is contained in:
Prajwal Raymond Moras 2026-02-23 17:00:09 +05:30
parent 3561bfb96c
commit 9875b5b98e
7 changed files with 1443 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

388
cookbook/rag_gateway/README.md Normal file
View file

@ -0,0 +1,388 @@
# LiteLLM RAG Gateway Cookbook
A comprehensive example demonstrating production-ready RAG (Retrieval-Augmented Generation) systems using **LiteLLM** as a multi-provider LLM gateway, integrated with **LlamaIndex**, **Haystack**, and evaluated with **RAGAS**.
## 🎯 What This Example Demonstrates
This cookbook showcases:
- **Multi-Provider LLM Gateway**: Use LiteLLM to seamlessly switch between OpenAI, Anthropic, Groq, Ollama, and 100+ other providers
- **Production-Ready Patterns**: Automatic retries, fallback chains, timeouts, and comprehensive logging
- **RAG with LlamaIndex**: Build semantic search pipelines with vector embeddings and ChromaDB
- **RAG with Haystack**: Create modular, component-based RAG pipelines
- **Evaluation with RAGAS**: Measure faithfulness, answer relevancy, context precision, and recall
- **Provider Comparison**: A/B test different models and providers with consistent code
## 📁 Project Structure
```
rag_gateway/
├── README.md # This file
├── requirements.txt # Python dependencies
├── litellm_config.yaml # Multi-provider configuration (optional)
├── rag_llamaindex.py # LlamaIndex RAG pipeline
├── rag_haystack.py # Haystack RAG pipeline
├── evaluate_with_ragas.py # RAGAS evaluation script
└── sample_data/
├── documents.json # Knowledge base (10 AI/ML docs)
└── eval_dataset.json # Evaluation questions & ground truth
```
## 🚀 Quick Start
### 1. Prerequisites
- Python 3.9+ (tested with 3.12)
- API keys for LLM providers (at minimum, OpenAI)
### 2. Installation
```bash
# Navigate to this directory
cd cookbook/rag_gateway
# Create virtual environment (recommended)
python3 -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# Install dependencies
pip install -r requirements.txt
```
### 3. Set API Keys
```bash
# Required for most examples
export OPENAI_API_KEY="your-openai-key-here"
# Optional: For testing alternative providers
export ANTHROPIC_API_KEY="your-anthropic-key-here"
export GROQ_API_KEY="your-groq-key-here"
# For local models (free, no API key needed)
# Install Ollama: https://ollama.ai
# ollama pull llama3
```
### 4. Run Examples
#### LlamaIndex RAG Pipeline
```bash
python rag_llamaindex.py
```
**What it does:**
- Loads 10 AI/ML documents into memory
- Creates vector embeddings using HuggingFace models
- Stores embeddings in ChromaDB (in-memory)
- Runs 3 example queries with semantic search
- Saves results to `llamaindex_results.json`
**Expected output:**
```
======================================================================
LlamaIndex RAG with LiteLLM Multi-Provider Gateway
======================================================================
✅ Loaded 10 documents
✅ Configured LiteLLM with model: gpt-3.5-turbo
✅ Configured HuggingFace embeddings
🔄 Creating vector index (this may take a moment)...
✅ RAG pipeline created successfully
❓ Question: What are the main components of RAG architecture?
📚 Retrieved Sources:
1. RAG Architecture and Benefits (score: 0.892)
Preview: Retrieval-Augmented Generation (RAG) combines...
💡 Answer: The main components of RAG architecture are...
```
#### Haystack RAG Pipeline
```bash
python rag_haystack.py
```
**What it does:**
- Demonstrates Haystack's component-based pipeline architecture
- Uses custom LiteLLM generator component
- Shows how to integrate LiteLLM with Haystack 2.0
- Saves results to `haystack_results.json`
#### RAGAS Evaluation
```bash
python evaluate_with_ragas.py
```
**What it does:**
- Loads evaluation dataset with 8 questions
- Generates answers using RAG pipeline
- Evaluates with RAGAS metrics (faithfulness, relevancy, precision, recall)
- Compares multiple models/providers
- Saves detailed results to `ragas_evaluation_results.json`
**Expected output:**
```
======================================================================
RAGAS Evaluation Comparison
======================================================================
Model Faith Relev Prec Recall Overall
----------------------------------------------------------------------
gpt-3.5-turbo 0.892 0.945 0.878 0.901 0.904
gpt-4-turbo 0.934 0.967 0.912 0.945 0.940
======================================================================
```
## 🔧 Switching LLM Providers
One of LiteLLM's key benefits is **zero-code provider switching**. Simply change the `model_name` parameter:
### In Python Scripts
```python
# OpenAI
llm = setup_litellm(model_name="gpt-3.5-turbo")
llm = setup_litellm(model_name="gpt-4-turbo")
# Anthropic Claude
llm = setup_litellm(model_name="claude-3-haiku-20240307")
llm = setup_litellm(model_name="claude-3-sonnet-20240229")
# Groq (ultra-fast inference)
llm = setup_litellm(model_name="groq/llama3-70b-8192")
llm = setup_litellm(model_name="groq/mixtral-8x7b-32768")
# Local Ollama (free, private)
llm = setup_litellm(model_name="ollama/llama3")
llm = setup_litellm(model_name="ollama/mistral")
# Azure OpenAI
llm = setup_litellm(model_name="azure/gpt-4-deployment-name")
```
### Using Configuration File (Optional)
The `litellm_config.yaml` file demonstrates advanced configuration:
```yaml
model_list:
- model_name: gpt-4-turbo
litellm_params:
model: gpt-4-turbo-preview
api_key: os.environ/OPENAI_API_KEY
timeout: 30
max_retries: 3
router_settings:
routing_strategy: simple-shuffle
num_retries: 3
fallbacks:
- gpt-4-turbo
- gpt-3.5-turbo
- claude-3-haiku
```
## 📊 Understanding RAGAS Metrics
RAGAS evaluates RAG systems using LLM-based metrics:
| Metric | What It Measures | Good Score |
|--------|------------------|------------|
| **Faithfulness** | Is the answer grounded in retrieved context? (hallucination check) | > 0.8 |
| **Answer Relevancy** | Does the answer address the question? | > 0.8 |
| **Context Precision** | Are retrieved documents relevant to the question? | > 0.7 |
| **Context Recall** | Was all necessary information retrieved? | > 0.7 |
Higher scores are better. RAGAS uses an LLM as a judge, making evaluation more aligned with human judgment than traditional metrics like BLEU or ROUGE.
## 🎓 Key Concepts
### Why LiteLLM?
1. **Unified Interface**: One API for 100+ providers
2. **Reliability**: Built-in retries, fallbacks, timeouts
3. **Cost Tracking**: Monitor spending per request
4. **Easy Testing**: Switch providers without code changes
5. **Production Ready**: Used by companies for high-availability systems
### RAG Architecture
```
User Query
↓
[Embedding Model] → Query Vector
↓
[Vector Database] → Retrieve Top-K Documents
↓
[LLM via LiteLLM] → Generate Answer from Context
↓
Answer + Sources
```
### Production Best Practices Demonstrated
- ✅ **Retry Logic**: Automatic retries with exponential backoff
- ✅ **Fallback Chains**: Switch to backup providers on failure
- ✅ **Timeouts**: Prevent hanging requests
- ✅ **Logging**: Track all requests for debugging
- ✅ **Cost Tracking**: Monitor LLM API costs
- ✅ **Evaluation**: Measure quality with RAGAS
- ✅ **Source Attribution**: Return retrieved documents
- ✅ **Error Handling**: Graceful degradation
## 🔍 Sample Data
The `sample_data/` directory contains:
### documents.json
10 comprehensive documents covering:
- Large Language Models
- RAG Architecture
- LiteLLM Features
- Vector Embeddings
- LlamaIndex & Haystack
- RAGAS Evaluation
- Production Best Practices
- Prompt Engineering
- Cost Optimization
### eval_dataset.json
8 evaluation questions with:
- Question text
- Ground truth answer
- Expected source document IDs
Perfect for testing and benchmarking your RAG system.
## 🛠️ Customization
### Add Your Own Documents
Edit `sample_data/documents.json`:
```json
[
{
"id": "doc_11",
"title": "Your Document Title",
"content": "Your document content here..."
}
]
```
### Change Embedding Models
In the scripts, modify:
```python
# Use different HuggingFace model
embed_model = HuggingFaceEmbedding(
model_name="BAAI/bge-large-en-v1.5" # Larger, more accurate
)
# Or use OpenAI embeddings
from llama_index.embeddings.openai import OpenAIEmbedding
embed_model = OpenAIEmbedding(model="text-embedding-3-large")
```
### Adjust Retrieval Parameters
```python
# Retrieve more documents
query_engine = index.as_query_engine(
similarity_top_k=5, # Default is 3
response_mode="tree_summarize" # Different synthesis strategy
)
```
### Add More Models to Compare
In `evaluate_with_ragas.py`:
```python
models_to_test = [
"gpt-3.5-turbo",
"gpt-4-turbo",
"claude-3-haiku-20240307",
"claude-3-sonnet-20240229",
"groq/llama3-70b-8192",
"ollama/llama3",
]
```
## 🐛 Troubleshooting
### "API key not found"
```bash
# Make sure you've exported your API key
echo $OPENAI_API_KEY
# If empty, set it:
export OPENAI_API_KEY="sk-..."
```
### "Module not found" errors
```bash
# Reinstall dependencies
pip install -r requirements.txt --upgrade
```
### Slow embedding generation
First run downloads HuggingFace models (~100MB). Subsequent runs are fast.
### RAGAS evaluation fails
RAGAS requires an LLM for evaluation. Ensure `OPENAI_API_KEY` is set.
### Ollama connection refused
```bash
# Start Ollama server
ollama serve
# In another terminal, pull a model
ollama pull llama3
```
## 📚 Learn More
- **LiteLLM Docs**: https://docs.litellm.ai/
- **LlamaIndex Docs**: https://docs.llamaindex.ai/
- **Haystack Docs**: https://docs.haystack.deepset.ai/
- **RAGAS Docs**: https://docs.ragas.io/
## 🤝 Contributing
This example is part of the LiteLLM project. To contribute:
1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Submit a pull request
## 📝 License
This example follows the LiteLLM project license.
## 🎉 What's Next?
After running this example, you can:
1. **Deploy to Production**: Use LiteLLM Proxy for authentication, rate limiting, and load balancing
2. **Scale Up**: Replace in-memory vector stores with Pinecone, Weaviate, or Qdrant
3. **Add Caching**: Implement semantic caching to reduce API costs
4. **Fine-tune Retrieval**: Experiment with hybrid search (semantic + keyword)
5. **Monitor Performance**: Integrate with Langfuse or other observability tools
6. **Build an API**: Wrap your RAG pipeline in FastAPI or Flask
7. **Add Streaming**: Enable streaming responses for better UX
---
**Built with ❤️ to demonstrate LiteLLM's power for production RAG systems**

352
cookbook/rag_gateway/evaluate_with_ragas.py Normal file
View file

@ -0,0 +1,352 @@
"""
RAG Evaluation with RAGAS and LiteLLM
This script demonstrates:
- Evaluating RAG systems with RAGAS metrics
- Comparing different LLM providers/configurations
- Measuring faithfulness, answer relevancy, and context precision
- Using LiteLLM for evaluation LLM calls
"""
import os
import json
from pathlib import Path
from typing import List, Dict
import warnings
warnings.filterwarnings('ignore')
from datasets import Dataset
from ragas import evaluate
from ragas.metrics import (
faithfulness,
answer_relevancy,
context_precision,
context_recall,
)
from llama_index.core import VectorStoreIndex, Document, Settings
from llama_index.llms.litellm import LiteLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
def load_documents(data_path: str = "sample_data/documents.json") -> List[Document]:
"""Load documents from JSON file."""
script_dir = Path(__file__).parent
file_path = script_dir / data_path
with open(file_path, 'r') as f:
docs_data = json.load(f)
documents = []
for doc_data in docs_data:
doc = Document(
text=doc_data['content'],
metadata={'doc_id': doc_data['id'], 'title': doc_data['title']}
)
documents.append(doc)
return documents
def load_eval_dataset(data_path: str = "sample_data/eval_dataset.json") -> List[Dict]:
"""Load evaluation dataset."""
script_dir = Path(__file__).parent
file_path = script_dir / data_path
with open(file_path, 'r') as f:
eval_data = json.load(f)
print(f"✅ Loaded {len(eval_data)} evaluation questions")
return eval_data
def create_rag_index(
documents: List[Document],
model_name: str,
collection_name: str = "ragas_eval"
) -> VectorStoreIndex:
"""Create RAG index with specified model."""
llm = LiteLLM(
model=model_name,
temperature=0.1,
max_tokens=512,
num_retries=3,
timeout=30.0,
)
embed_model = HuggingFaceEmbedding(
model_name="BAAI/bge-small-en-v1.5",
pooling="cls",
)
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512
Settings.chunk_overlap = 50
# Create vector store
chroma_client = chromadb.EphemeralClient()
chroma_collection = chroma_client.create_collection(collection_name)
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=False,
)
return index
def generate_rag_responses(
index: VectorStoreIndex,
eval_dataset: List[Dict],
top_k: int = 3
) -> Dict[str, List]:
"""
Generate RAG responses for evaluation dataset.
Returns data in RAGAS format:
- question: User question
- answer: Generated answer
- contexts: Retrieved document chunks
- ground_truth: Expected answer
"""
questions = []
answers = []
contexts = []
ground_truths = []
query_engine = index.as_query_engine(
similarity_top_k=top_k,
response_mode="compact",
)
print(f"🔄 Generating responses for {len(eval_dataset)} questions...")
for i, item in enumerate(eval_dataset, 1):
question = item['question']
ground_truth = item['ground_truth']
# Query RAG system
response = query_engine.query(question)
answer = str(response.response)
# Extract contexts
retrieved_contexts = []
if hasattr(response, 'source_nodes'):
for node in response.source_nodes:
retrieved_contexts.append(node.text)
questions.append(question)
answers.append(answer)
contexts.append(retrieved_contexts)
ground_truths.append(ground_truth)
print(f" ✓ Processed {i}/{len(eval_dataset)}: {question[:60]}...")
return {
'question': questions,
'answer': answers,
'contexts': contexts,
'ground_truth': ground_truths,
}
def evaluate_rag_system(
rag_data: Dict[str, List],
model_name: str,
eval_model: str = "gpt-3.5-turbo"
) -> Dict:
"""
Evaluate RAG system using RAGAS metrics.
Metrics:
- Faithfulness: Is answer grounded in retrieved context?
- Answer Relevancy: Does answer address the question?
- Context Precision: Are retrieved docs relevant?
- Context Recall: Was all needed info retrieved?
"""
print(f"\n📊 Evaluating with RAGAS (using {eval_model} as judge)...")
# Convert to RAGAS dataset format
dataset = Dataset.from_dict(rag_data)
# Configure RAGAS to use LiteLLM
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "dummy-key")
# Run evaluation
try:
result = evaluate(
dataset,
metrics=[
faithfulness,
answer_relevancy,
context_precision,
context_recall,
],
)
scores = {
'model': model_name,
'faithfulness': result['faithfulness'],
'answer_relevancy': result['answer_relevancy'],
'context_precision': result['context_precision'],
'context_recall': result['context_recall'],
'overall_score': (
result['faithfulness'] +
result['answer_relevancy'] +
result['context_precision'] +
result['context_recall']
) / 4
}
return scores
except Exception as e:
print(f"❌ Evaluation error: {e}")
return {
'model': model_name,
'error': str(e),
'faithfulness': 0.0,
'answer_relevancy': 0.0,
'context_precision': 0.0,
'context_recall': 0.0,
'overall_score': 0.0,
}
def compare_models(
documents: List[Document],
eval_dataset: List[Dict],
models: List[str]
) -> List[Dict]:
"""
Compare multiple models/providers using RAGAS evaluation.
This demonstrates how LiteLLM enables easy A/B testing
across different providers and model configurations.
"""
results = []
for model in models:
print("\n" + "=" * 70)
print(f"Evaluating Model: {model}")
print("=" * 70)
try:
# Create RAG index with this model
index = create_rag_index(
documents,
model_name=model,
collection_name=f"eval_{model.replace('/', '_')}"
)
# Generate responses
rag_data = generate_rag_responses(index, eval_dataset, top_k=3)
# Evaluate
scores = evaluate_rag_system(rag_data, model_name=model)
results.append(scores)
# Print results
print(f"\n✅ Results for {model}:")
print(f" Faithfulness: {scores['faithfulness']:.3f}")
print(f" Answer Relevancy: {scores['answer_relevancy']:.3f}")
print(f" Context Precision: {scores['context_precision']:.3f}")
print(f" Context Recall: {scores['context_recall']:.3f}")
print(f" Overall Score: {scores['overall_score']:.3f}")
except Exception as e:
print(f"❌ Error evaluating {model}: {e}")
results.append({
'model': model,
'error': str(e),
'overall_score': 0.0
})
return results
def print_comparison_table(results: List[Dict]):
"""Print formatted comparison table."""
print("\n" + "=" * 70)
print("RAGAS Evaluation Comparison")
print("=" * 70)
print(f"{'Model':<30} {'Faith':<8} {'Relev':<8} {'Prec':<8} {'Recall':<8} {'Overall':<8}")
print("-" * 70)
for result in results:
if 'error' not in result:
print(f"{result['model']:<30} "
f"{result['faithfulness']:<8.3f} "
f"{result['answer_relevancy']:<8.3f} "
f"{result['context_precision']:<8.3f} "
f"{result['context_recall']:<8.3f} "
f"{result['overall_score']:<8.3f}")
else:
print(f"{result['model']:<30} ERROR: {result['error']}")
print("=" * 70)
def main():
"""Main execution flow."""
print("=" * 70)
print("RAG Evaluation with RAGAS and LiteLLM")
print("=" * 70 + "\n")
# Check for API keys
if not os.getenv("OPENAI_API_KEY"):
print("⚠️ Warning: OPENAI_API_KEY not set.")
print(" RAGAS requires an LLM for evaluation metrics.")
print(" Set it with: export OPENAI_API_KEY='your-key-here'\n")
return
# Load data
documents = load_documents()
eval_dataset = load_eval_dataset()
# Models to compare
# Uncomment models you have API keys for
models_to_test = [
"gpt-3.5-turbo",
# "gpt-4-turbo",
# "claude-3-haiku-20240307",
# "groq/llama3-70b-8192",
]
print(f"\n🔬 Testing {len(models_to_test)} model configurations...")
# Run comparison
results = compare_models(documents, eval_dataset, models_to_test)
# Print comparison
print_comparison_table(results)
# Save results
output_path = Path(__file__).parent / "ragas_evaluation_results.json"
with open(output_path, 'w') as f:
json.dump(results, f, indent=2)
print(f"\n✅ Detailed results saved to: {output_path}")
# Insights
print("\n💡 Key Insights:")
print(" - RAGAS uses LLMs as judges for evaluation")
print(" - Faithfulness measures hallucination (higher is better)")
print(" - Answer Relevancy checks if answer addresses question")
print(" - Context metrics evaluate retrieval quality")
print(" - LiteLLM makes it easy to compare providers")
print(" - Add more models to models_to_test list for comparison")
print("\n" + "=" * 70)
print("🎉 Evaluation completed successfully!")
print("=" * 70)
if __name__ == "__main__":
main()

328
cookbook/rag_gateway/rag_haystack.py Normal file
View file

@ -0,0 +1,328 @@
"""
Haystack RAG Pipeline with LiteLLM Multi-Provider Gateway
This script demonstrates:
- Building a RAG pipeline with Haystack 2.0
- Using LiteLLM for flexible LLM provider switching
- Document ingestion and retrieval
- Answer generation with source attribution
"""
import os
import json
from pathlib import Path
from typing import List, Dict, Any
from haystack import Pipeline, Document
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.builders import PromptBuilder
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
from haystack.components.retrievers import InMemoryEmbeddingRetriever
import litellm
def load_documents(data_path: str = "sample_data/documents.json") -> List[Document]:
"""Load documents from JSON and convert to Haystack format."""
script_dir = Path(__file__).parent
file_path = script_dir / data_path
with open(file_path, 'r') as f:
docs_data = json.load(f)
documents = []
for doc_data in docs_data:
doc = Document(
content=doc_data['content'],
meta={
'doc_id': doc_data['id'],
'title': doc_data['title']
}
)
documents.append(doc)
print(f"✅ Loaded {len(documents)} documents")
return documents
class LiteLLMGenerator:
"""
Custom Haystack component that uses LiteLLM for generation.
This allows seamless switching between providers (OpenAI, Anthropic, etc.)
while maintaining Haystack's pipeline architecture.
"""
def __init__(
self,
model: str = "gpt-3.5-turbo",
temperature: float = 0.1,
max_tokens: int = 512,
num_retries: int = 3,
timeout: float = 30.0
):
self.model = model
self.temperature = temperature
self.max_tokens = max_tokens
self.num_retries = num_retries
self.timeout = timeout
# Configure LiteLLM settings
litellm.num_retries = num_retries
litellm.request_timeout = timeout
print(f"✅ Configured LiteLLM Generator with model: {model}")
def run(self, prompt: str) -> Dict[str, Any]:
"""
Generate response using LiteLLM.
LiteLLM handles:
- Provider-specific API translation
- Automatic retries with exponential backoff
- Fallback to alternative providers
- Cost tracking and logging
"""
try:
response = litellm.completion(
model=self.model,
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
max_tokens=self.max_tokens,
num_retries=self.num_retries,
timeout=self.timeout,
stream=False, # Explicitly disable streaming
# Optional: Enable fallbacks
fallbacks=["gpt-3.5-turbo", "claude-3-haiku-20240307"],
)
# Safely extract answer
answer = ""
if hasattr(response, 'choices') and len(response.choices) > 0: # type: ignore
answer = response.choices[0].message.content # type: ignore
else:
answer = str(response)
# Safely extract metadata
meta: Dict[str, Any] = {"model": self.model}
if hasattr(response, 'model') and response.model:
meta["model"] = str(response.model)
if hasattr(response, 'usage') and response.usage: # type: ignore
meta["usage"] = {
"prompt_tokens": int(response.usage.prompt_tokens), # type: ignore
"completion_tokens": int(response.usage.completion_tokens), # type: ignore
"total_tokens": int(response.usage.total_tokens), # type: ignore
}
return {
"replies": [answer],
"meta": meta
}
except Exception as e:
print(f"❌ Error in LiteLLM generation: {e}")
return {
"replies": [f"Error: {str(e)}"],
"meta": {"error": str(e)}
}
def create_document_store_with_embeddings(documents: List[Document]) -> InMemoryDocumentStore:
"""
Create document store and add embeddings for semantic search.
"""
# Initialize document store
document_store = InMemoryDocumentStore()
# Create embedder
doc_embedder = SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2"
)
doc_embedder.warm_up()
# Generate embeddings
print("🔄 Generating embeddings for documents...")
docs_with_embeddings = doc_embedder.run(documents)
# Write to store
document_store.write_documents(docs_with_embeddings["documents"])
print(f"✅ Document store created with {document_store.count_documents()} documents")
return document_store
def create_rag_pipeline(
document_store: InMemoryDocumentStore,
model_name: str = "gpt-3.5-turbo"
) -> Pipeline:
"""
Create Haystack RAG pipeline with LiteLLM generator.
Pipeline components:
1. Text Embedder - Convert query to embedding
2. Retriever - Fetch relevant documents
3. Prompt Builder - Format context and question
4. LiteLLM Generator - Generate answer
"""
# Initialize components
text_embedder = SentenceTransformersTextEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2"
)
retriever = InMemoryEmbeddingRetriever(document_store=document_store)
# RAG prompt template
template = """
You are a helpful AI assistant. Answer the question based on the provided context.
If the context doesn't contain enough information, say so.
Context:
{% for document in documents %}
{{ document.content }}
{% endfor %}
Question: {{ question }}
Answer:
"""
prompt_builder = PromptBuilder(template=template)
llm_generator = LiteLLMGenerator(model=model_name)
# Build pipeline
pipeline = Pipeline()
pipeline.add_component("text_embedder", text_embedder)
pipeline.add_component("retriever", retriever)
pipeline.add_component("prompt_builder", prompt_builder)
pipeline.add_component("llm", llm_generator)
# Connect components
pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
pipeline.connect("retriever.documents", "prompt_builder.documents")
pipeline.connect("prompt_builder.prompt", "llm.prompt")
print("✅ Haystack RAG pipeline created successfully")
return pipeline
def query_rag(
pipeline: Pipeline,
question: str,
top_k: int = 3,
verbose: bool = True
) -> Dict:
"""
Query the RAG pipeline and return answer with sources.
Args:
pipeline: Haystack pipeline
question: User question
top_k: Number of documents to retrieve
verbose: Print retrieved context
Returns:
Dict with answer, sources, and metadata
"""
print(f"\n❓ Question: {question}")
# Run pipeline
result = pipeline.run({
"text_embedder": {"text": question},
"retriever": {"top_k": top_k},
"prompt_builder": {"question": question}
})
# Extract answer
answer = result["llm"]["replies"][0]
metadata = result["llm"].get("meta", {})
# Extract sources
sources = []
if "retriever" in result and "documents" in result["retriever"]:
for doc in result["retriever"]["documents"]:
sources.append({
'doc_id': doc.meta.get('doc_id', 'unknown'),
'title': doc.meta.get('title', 'unknown'),
'score': doc.score if hasattr(doc, 'score') else 0.0,
'text_preview': doc.content[:200] + "..."
})
if verbose and sources:
print("\n📚 Retrieved Sources:")
for i, source in enumerate(sources, 1):
print(f" {i}. {source['title']} (score: {source['score']:.3f})")
print(f" Preview: {source['text_preview']}\n")
print(f"💡 Answer: {answer}\n")
if verbose and metadata.get('usage'):
usage = metadata['usage']
print(f"📊 Token Usage: {usage['total_tokens']} total "
f"({usage['prompt_tokens']} prompt + {usage['completion_tokens']} completion)")
return {
'question': question,
'answer': answer,
'sources': sources,
'metadata': metadata
}
def main():
"""Main execution flow."""
print("=" * 70)
print("Haystack RAG with LiteLLM Multi-Provider Gateway")
print("=" * 70 + "\n")
# Check for API keys
if not os.getenv("OPENAI_API_KEY"):
print("⚠️ Warning: OPENAI_API_KEY not set. Set it with:")
print(" export OPENAI_API_KEY='your-key-here'\n")
print(" Attempting to use fallback providers...\n")
# Load documents
documents = load_documents()
# Create document store with embeddings
document_store = create_document_store_with_embeddings(documents)
# Create RAG pipeline
# Try different models: "gpt-4-turbo", "gpt-3.5-turbo",
# "claude-3-haiku-20240307", "groq/llama3-70b-8192"
pipeline = create_rag_pipeline(document_store, model_name="gpt-3.5-turbo")
# Example queries
questions = [
"What are the key differences between LlamaIndex and Haystack?",
"How does LiteLLM provide reliability for production systems?",
"What are best practices for chunking documents in RAG?",
]
results = []
for question in questions:
result = query_rag(pipeline, question, top_k=3)
results.append(result)
print("-" * 70 + "\n")
# Save results
output_path = Path(__file__).parent / "haystack_results.json"
with open(output_path, 'w') as f:
json.dump(results, f, indent=2)
print(f"✅ Results saved to: {output_path}")
print("\n" + "=" * 70)
print("🎉 Haystack RAG pipeline completed successfully!")
print("=" * 70)
# Tips
print("\n💡 Tips:")
print(" - Haystack's component architecture makes pipelines modular")
print(" - Switch LLM providers by changing model_name parameter")
print(" - LiteLLM handles retries, fallbacks, and cost tracking")
print(" - Combine with BM25 retriever for hybrid search")
print(" - Add re-rankers for improved retrieval quality")
if __name__ == "__main__":
main()

256
cookbook/rag_gateway/rag_llamaindex.py Normal file
View file

@ -0,0 +1,256 @@
"""
LlamaIndex RAG Pipeline with LiteLLM Multi-Provider Gateway
This script demonstrates:
- Building a RAG pipeline with LlamaIndex
- Using LiteLLM for multi-provider LLM access
- Configuring retries, fallbacks, and timeouts
- Vector-based semantic search with ChromaDB
"""
import os
import json
from pathlib import Path
from typing import List, Dict
from llama_index.core import (
VectorStoreIndex,
Document,
Settings,
StorageContext,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.litellm import LiteLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
def load_documents(data_path: str = "sample_data/documents.json") -> List[Document]:
"""Load documents from JSON file."""
script_dir = Path(__file__).parent
file_path = script_dir / data_path
with open(file_path, 'r') as f:
docs_data = json.load(f)
documents = []
for doc_data in docs_data:
doc = Document(
text=doc_data['content'],
metadata={
'doc_id': doc_data['id'],
'title': doc_data['title']
}
)
documents.append(doc)
print(f"✅ Loaded {len(documents)} documents")
return documents
def setup_litellm(
model_name: str = "gpt-3.5-turbo",
temperature: float = 0.1,
max_tokens: int = 512
) -> LiteLLM:
"""
Configure LiteLLM with retry and fallback settings.
LiteLLM automatically handles:
- Provider-specific API format translation
- Retries with exponential backoff
- Fallback to alternative providers
- Cost tracking and logging
"""
llm = LiteLLM(
model=model_name,
temperature=temperature,
max_tokens=max_tokens,
# These settings enable production-ready reliability
num_retries=3, # Retry failed requests
timeout=30.0, # 30 second timeout
# Fallback models (if primary fails)
fallbacks=["gpt-3.5-turbo", "claude-3-haiku-20240307"],
)
print(f"✅ Configured LiteLLM with model: {model_name}")
return llm
def setup_embeddings() -> HuggingFaceEmbedding:
"""Configure embedding model for semantic search."""
embed_model = HuggingFaceEmbedding(
model_name="BAAI/bge-small-en-v1.5",
# Use 'cls' pooling for better performance
pooling="cls",
)
print("✅ Configured HuggingFace embeddings")
return embed_model
def create_rag_pipeline(
documents: List[Document],
llm: LiteLLM,
embed_model: HuggingFaceEmbedding,
collection_name: str = "litellm_rag_docs"
) -> VectorStoreIndex:
"""
Create RAG pipeline with vector store and query engine.
Pipeline steps:
1. Chunk documents into smaller pieces
2. Generate embeddings for each chunk
3. Store in ChromaDB vector database
4. Create query engine for retrieval + generation
"""
# Configure global settings
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512
Settings.chunk_overlap = 50
# Initialize ChromaDB
chroma_client = chromadb.EphemeralClient()
chroma_collection = chroma_client.create_collection(collection_name)
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# Parse documents into nodes (chunks)
node_parser = SentenceSplitter(
chunk_size=512,
chunk_overlap=50,
)
# Create index with embeddings
print("🔄 Creating vector index (this may take a moment)...")
index = VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
node_parser=node_parser,
show_progress=True,
)
print("✅ RAG pipeline created successfully")
return index
def query_rag(
index: VectorStoreIndex,
question: str,
top_k: int = 3,
verbose: bool = True
) -> Dict:
"""
Query the RAG system and return answer with sources.
Args:
index: Vector store index
question: User question
top_k: Number of documents to retrieve
verbose: Print retrieved context
Returns:
Dict with answer, sources, and metadata
"""
# Create query engine with custom prompt
query_engine = index.as_query_engine(
similarity_top_k=top_k,
response_mode="compact", # Combine context efficiently
)
# Execute query
print(f"\n❓ Question: {question}")
response = query_engine.query(question)
# Extract source information
sources = []
if hasattr(response, 'source_nodes'):
for node in response.source_nodes:
sources.append({
'doc_id': node.metadata.get('doc_id', 'unknown'),
'title': node.metadata.get('title', 'unknown'),
'score': node.score,
'text_preview': node.text[:200] + "..."
})
if verbose and sources:
print("\n📚 Retrieved Sources:")
for i, source in enumerate(sources, 1):
print(f" {i}. {source['title']} (score: {source['score']:.3f})")
print(f" Preview: {source['text_preview']}\n")
print(f"💡 Answer: {response.response}\n")
return {
'question': question,
'answer': str(response.response),
'sources': sources,
'metadata': {
'model': Settings.llm.model,
'top_k': top_k
}
}
def main():
"""Main execution flow."""
print("=" * 70)
print("LlamaIndex RAG with LiteLLM Multi-Provider Gateway")
print("=" * 70 + "\n")
# Check for API keys
if not os.getenv("OPENAI_API_KEY"):
print("⚠️ Warning: OPENAI_API_KEY not set. Set it with:")
print(" export OPENAI_API_KEY='your-key-here'\n")
print(" Attempting to use fallback providers...\n")
# Load documents
documents = load_documents()
# Setup LiteLLM (try different models by changing model_name)
# Options: "gpt-4-turbo", "gpt-3.5-turbo", "claude-3-haiku-20240307",
# "groq/llama3-70b-8192", "ollama/llama3"
llm = setup_litellm(model_name="gpt-3.5-turbo")
# Setup embeddings
embed_model = setup_embeddings()
# Create RAG pipeline
index = create_rag_pipeline(documents, llm, embed_model)
# Example queries
questions = [
"What are the main components of RAG architecture?",
"How does LiteLLM help with multi-provider integration?",
"What metrics does RAGAS use for evaluation?",
]
results = []
for question in questions:
result = query_rag(index, question, top_k=3)
results.append(result)
print("-" * 70 + "\n")
# Save results
output_path = Path(__file__).parent / "llamaindex_results.json"
with open(output_path, 'w') as f:
json.dump(results, f, indent=2)
print(f"✅ Results saved to: {output_path}")
print("\n" + "=" * 70)
print("🎉 LlamaIndex RAG pipeline completed successfully!")
print("=" * 70)
# Tips for switching providers
print("\n💡 Tips:")
print(" - Change model by modifying 'model_name' parameter")
print(" - Set ANTHROPIC_API_KEY to use Claude models")
print(" - Set GROQ_API_KEY for ultra-fast inference")
print(" - Install Ollama for free local models")
print(" - LiteLLM automatically handles retries and fallbacks")
if __name__ == "__main__":
main()

25
cookbook/rag_gateway/requirements.txt Normal file
View file

@ -0,0 +1,25 @@
# Core LiteLLM
litellm>=1.50.0
# RAG Frameworks
llama-index>=0.10.0
llama-index-llms-litellm>=0.1.0
llama-index-embeddings-huggingface>=0.2.0
haystack-ai>=2.0.0
# Evaluation
ragas>=0.1.0
datasets>=2.14.0
# Vector Store & Embeddings
chromadb>=0.4.0
sentence-transformers>=2.2.0
# Utilities
python-dotenv>=1.0.0
pyyaml>=6.0
# Use --only-binary to avoid build issues on Python 3.15
# Install these separately if needed:
# pip install numpy --only-binary :all:
# pip install pandas --only-binary :all:

52
cookbook/rag_gateway/sample_data/documents.json Normal file
View file

@ -0,0 +1,52 @@
[
{
"id": "doc_1",
"title": "Introduction to Large Language Models",
"content": "Large Language Models (LLMs) are neural networks trained on vast amounts of text data. They use transformer architecture with self-attention mechanisms to understand and generate human-like text. Popular LLMs include GPT-4, Claude, and Llama. These models are pre-trained on diverse internet text and can be fine-tuned for specific tasks. LLMs have revolutionized natural language processing by achieving state-of-the-art results on tasks like translation, summarization, question answering, and code generation. The training process involves predicting the next token in a sequence, which allows the model to learn patterns, grammar, facts, and reasoning abilities from the training data."
},
{
"id": "doc_2",
"title": "RAG Architecture and Benefits",
"content": "Retrieval-Augmented Generation (RAG) combines information retrieval with language generation. The architecture has three main components: a retriever that fetches relevant documents from a knowledge base, an embedder that converts text to vector representations, and a generator (LLM) that produces answers using retrieved context. RAG reduces hallucinations by grounding responses in factual documents, enables knowledge updates without retraining, and provides source attribution for transparency. It's particularly effective for domain-specific applications where up-to-date information is critical. RAG systems typically use vector databases like ChromaDB, Pinecone, or Weaviate to store and retrieve document embeddings efficiently. The retrieval step uses semantic similarity to find the most relevant documents for a given query."
},
{
"id": "doc_3",
"title": "LiteLLM Multi-Provider Gateway",
"content": "LiteLLM is a unified interface for 100+ LLM providers including OpenAI, Anthropic, Azure, AWS Bedrock, Google Vertex AI, and local models like Ollama. It translates inputs to provider-specific formats and returns consistent OpenAI-compatible outputs. Key features include automatic retries with exponential backoff, fallback chains across providers for reliability, request timeouts to prevent hanging, cost tracking per request, and comprehensive logging for observability. LiteLLM's router enables load balancing across multiple model deployments, making it ideal for production RAG systems that need high availability. The proxy server adds authentication, rate limiting, and budget management capabilities. Configuration is done via YAML files that define model lists, retry policies, and fallback strategies."
},
{
"id": "doc_4",
"title": "Vector Embeddings and Semantic Search",
"content": "Vector embeddings are dense numerical representations of text that capture semantic meaning. Words or sentences with similar meanings have embeddings that are close together in high-dimensional space. Embedding models like sentence-transformers, OpenAI's text-embedding-ada-002, or Cohere's embed models convert text into vectors (typically 384 to 1536 dimensions). Semantic search uses these embeddings to find relevant documents by computing cosine similarity or dot product between query and document vectors. This approach outperforms traditional keyword-based search because it understands context and meaning. For RAG systems, embeddings are pre-computed for all documents and stored in a vector database with efficient indexing (HNSW, IVF) for fast retrieval. At query time, the question is embedded and used to retrieve the top-k most similar documents."
},
{
"id": "doc_5",
"title": "LlamaIndex Framework Overview",
"content": "LlamaIndex is a data framework for building LLM applications with external data. It provides data connectors to ingest from APIs, databases, PDFs, and other sources. The core abstraction is the Index, which structures data for efficient retrieval. VectorStoreIndex creates embeddings and enables semantic search. ListIndex stores documents sequentially. TreeIndex builds a hierarchical structure. LlamaIndex includes a query engine that orchestrates retrieval and generation, supporting various retrieval strategies like top-k similarity, MMR (Maximum Marginal Relevance), and hybrid search. It integrates with LiteLLM through the LiteLLM LLM class, allowing seamless provider switching. Response synthesizers combine retrieved context with LLM generation using techniques like refine, compact, or tree summarize. LlamaIndex also supports agents, tools, and multi-step reasoning workflows."
},
{
"id": "doc_6",
"title": "Haystack RAG Pipeline",
"content": "Haystack is an open-source framework for building production-ready RAG pipelines. It uses a pipeline architecture where components are connected as nodes in a directed graph. Key components include DocumentStores (Elasticsearch, Weaviate, Pinecone) for storing documents and embeddings, Retrievers (DenseRetriever, BM25Retriever) for fetching relevant documents, Readers/Generators for answer generation, and Rankers for re-ranking retrieved results. Haystack 2.0 introduced a more flexible component-based design with better typing and async support. It integrates with LiteLLM through custom generator components that call the LiteLLM completion API. Haystack excels at complex pipelines with multiple retrieval stages, question classification, and answer validation. It supports both extractive QA (extracting spans from documents) and generative QA (synthesizing answers with LLMs)."
},
{
"id": "doc_7",
"title": "RAGAS Evaluation Framework",
"content": "RAGAS (Retrieval-Augmented Generation Assessment) is a framework for evaluating RAG systems using LLM-based metrics. It measures both retrieval quality and generation quality. Key metrics include: Faithfulness (whether the answer is grounded in retrieved context), Answer Relevancy (how well the answer addresses the question), Context Precision (how relevant retrieved documents are), Context Recall (whether all necessary information was retrieved), and Answer Semantic Similarity (comparing generated answers to ground truth). RAGAS uses LLMs as judges to compute these metrics, making evaluation more aligned with human judgment than traditional metrics like BLEU or ROUGE. It requires a test dataset with questions, ground truth answers, retrieved contexts, and generated answers. RAGAS helps compare different retrieval strategies, chunk sizes, embedding models, and LLM providers to optimize RAG system performance."
},
{
"id": "doc_8",
"title": "Production RAG Best Practices",
"content": "Building production RAG systems requires careful attention to reliability, latency, and cost. Best practices include: implementing retry logic with exponential backoff for API failures, using fallback chains across multiple LLM providers for high availability, setting appropriate timeouts to prevent hanging requests, chunking documents optimally (typically 256-512 tokens with overlap), using hybrid search (combining semantic and keyword search) for better recall, implementing caching for frequently asked questions, monitoring retrieval quality metrics (MRR, NDCG), tracking LLM costs per request, adding guardrails to filter inappropriate content, and logging all requests for debugging and analysis. For scaling, consider async processing, batch embeddings, and distributed vector databases. Security measures include API key rotation, rate limiting per user, and input validation. Regular evaluation with RAGAS or similar frameworks ensures system quality over time."
},
{
"id": "doc_9",
"title": "Prompt Engineering for RAG",
"content": "Effective prompt engineering is crucial for RAG systems. The prompt should clearly instruct the LLM to use only the provided context, cite sources when possible, and admit when information is insufficient. A typical RAG prompt structure includes: system instructions defining the assistant's role, retrieved context documents with clear delimiters, the user's question, and output format instructions. Techniques like few-shot examples improve consistency. Chain-of-thought prompting helps with complex reasoning over multiple documents. For multi-turn conversations, include chat history while managing context window limits. Prompt compression techniques like selective context or summarization help fit more information. When using LiteLLM, prompts should be provider-agnostic since LiteLLM handles format translation. Testing prompts across different models (GPT-4, Claude, Llama) ensures robustness. Prompt versioning and A/B testing help optimize performance over time."
},
{
"id": "doc_10",
"title": "Cost Optimization for LLM Applications",
"content": "LLM costs can escalate quickly in production RAG systems. Optimization strategies include: using smaller models for simple queries and larger models only when needed, implementing semantic caching to avoid redundant API calls, compressing prompts by removing unnecessary context, batching requests when possible, using streaming for better user experience without cost increase, choosing cost-effective providers (Claude Haiku, GPT-3.5-turbo, or open-source models via Ollama), setting token limits to prevent runaway generation, monitoring costs per user/query with tools like LiteLLM's cost tracking, and implementing rate limiting and quotas. For embeddings, batch processing and caching reduce costs significantly. Consider fine-tuning smaller models for specific tasks instead of always using large general-purpose models. LiteLLM's router can automatically route to cheaper models based on query complexity. Regular cost analysis helps identify optimization opportunities."
}
]

42
cookbook/rag_gateway/sample_data/eval_dataset.json Normal file
View file

@ -0,0 +1,42 @@
[
{
"question": "What are the main components of RAG architecture?",
"ground_truth": "RAG architecture has three main components: a retriever that fetches relevant documents from a knowledge base, an embedder that converts text to vector representations, and a generator (LLM) that produces answers using retrieved context.",
"contexts": ["doc_2"]
},
{
"question": "How does LiteLLM help with multi-provider LLM integration?",
"ground_truth": "LiteLLM provides a unified interface for 100+ LLM providers, translates inputs to provider-specific formats, returns consistent OpenAI-compatible outputs, and includes features like automatic retries, fallback chains, timeouts, cost tracking, and comprehensive logging.",
"contexts": ["doc_3"]
},
{
"question": "What is the purpose of vector embeddings in semantic search?",
"ground_truth": "Vector embeddings are dense numerical representations that capture semantic meaning. They enable semantic search by allowing systems to find relevant documents through cosine similarity or dot product between query and document vectors, understanding context and meaning rather than just keywords.",
"contexts": ["doc_4"]
},
{
"question": "What retrieval strategies does LlamaIndex support?",
"ground_truth": "LlamaIndex supports various retrieval strategies including top-k similarity search, MMR (Maximum Marginal Relevance), and hybrid search that combines multiple approaches.",
"contexts": ["doc_5"]
},
{
"question": "What metrics does RAGAS use to evaluate RAG systems?",
"ground_truth": "RAGAS uses metrics including Faithfulness (answer grounded in context), Answer Relevancy (how well answer addresses question), Context Precision (relevance of retrieved documents), Context Recall (whether all necessary information was retrieved), and Answer Semantic Similarity (comparing to ground truth).",
"contexts": ["doc_7"]
},
{
"question": "What are best practices for production RAG systems?",
"ground_truth": "Best practices include implementing retry logic with exponential backoff, using fallback chains across providers, setting timeouts, optimal document chunking (256-512 tokens with overlap), hybrid search, caching, monitoring retrieval metrics, tracking costs, adding guardrails, and comprehensive logging.",
"contexts": ["doc_8"]
},
{
"question": "How can you optimize costs in LLM applications?",
"ground_truth": "Cost optimization strategies include using smaller models when appropriate, semantic caching, prompt compression, request batching, choosing cost-effective providers, setting token limits, monitoring per-query costs, rate limiting, and considering fine-tuned smaller models instead of large general-purpose ones.",
"contexts": ["doc_10"]
},
{
"question": "What are the key differences between LlamaIndex and Haystack?",
"ground_truth": "LlamaIndex focuses on data framework abstractions with various index types (VectorStoreIndex, ListIndex, TreeIndex) and response synthesizers. Haystack uses a pipeline architecture with components as nodes in a directed graph, excelling at complex multi-stage pipelines with question classification and answer validation.",
"contexts": ["doc_5", "doc_6"]
}
]