A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[doc] Add RAG Integration example (#17692)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
May 7, 2025 at 00:10 UTC
7525d5f3d52bcea2ada5af3205a6824ae62a737c
4 files changed
+551
docs/source/deployment/frameworks/index.md
+1
@@ -11,6 +11,7 @@ helm
11
lws
12
modal
13
open-webui
14
+retrieval_augmented_generation
15
skypilot
16
streamlit
17
triton
docs/source/deployment/frameworks/retrieval_augmented_generation.md
new
+84
@@ -0,0 +1,84 @@
1
+(deployment-retrieval-augmented-generation)=
2
+
3
+# Retrieval-Augmented Generation
4
+
5
+[Retrieval-augmented generation (RAG)](https://en.wikipedia.org/wiki/Retrieval-augmented_generation) is a technique that enables generative artificial intelligence (Gen AI) models to retrieve and incorporate new information. It modifies interactions with a large language model (LLM) so that the model responds to user queries with reference to a specified set of documents, using this information to supplement information from its pre-existing training data. This allows LLMs to use domain-specific and/or updated information. Use cases include providing chatbot access to internal company data or generating responses based on authoritative sources.
6
+
7
+Here are the integrations:
8
+- vLLM + [langchain](https://github.com/langchain-ai/langchain) + [milvus](https://github.com/milvus-io/milvus)
9
+- vLLM + [llamaindex](https://github.com/run-llama/llama_index) + [milvus](https://github.com/milvus-io/milvus)
10
+
11
+## vLLM + langchain
12
+
13
+### Prerequisites
14
+
15
+- Setup vLLM and langchain environment
16
+
17
+```console
18
+pip install -U vllm \
19
+ langchain_milvus langchain_openai \
20
+ langchain_community beautifulsoup4 \
21
+ langchain-text-splitters
22
+```
23
+
24
+### Deploy
25
+
26
+- Start the vLLM server with the supported embedding model, e.g.
27
+
28
+```console
29
+# Start embedding service (port 8000)
30
+vllm serve ssmits/Qwen2-7B-Instruct-embed-base
31
+```
32
+
33
+- Start the vLLM server with the supported chat completion model, e.g.
34
+
35
+```console
36
+# Start chat service (port 8001)
37
+vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
38
+```
39
+
40
+- Use the script: <gh-file:examples/online_serving/retrieval_augmented_generation_with_langchain.py>
41
+
42
+- Run the script
43
+
44
+```python
45
+python retrieval_augmented_generation_with_langchain.py
46
+```
47
+
48
+## vLLM + llamaindex
49
+
50
+### Prerequisites
51
+
52
+- Setup vLLM and llamaindex environment
53
+
54
+```console
55
+pip install vllm \
56
+ llama-index llama-index-readers-web \
57
+ llama-index-llms-openai-like \
58
+ llama-index-embeddings-openai-like \
59
+ llama-index-vector-stores-milvus \
60
+```
61
+
62
+### Deploy
63
+
64
+- Start the vLLM server with the supported embedding model, e.g.
65
+
66
+```console
67
+# Start embedding service (port 8000)
68
+vllm serve ssmits/Qwen2-7B-Instruct-embed-base
69
+```
70
+
71
+- Start the vLLM server with the supported chat completion model, e.g.
72
+
73
+```console
74
+# Start chat service (port 8001)
75
+vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
76
+```
77
+
78
+- Use the script: <gh-file:examples/online_serving/retrieval_augmented_generation_with_llamaindex.py>
79
+
80
+- Run the script
81
+
82
+```python
83
+python retrieval_augmented_generation_with_llamaindex.py
84
+```
examples/online_serving/retrieval_augmented_generation_with_langchain.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Retrieval Augmented Generation (RAG) Implementation with Langchain
4
+==================================================================
5
+
6
+This script demonstrates a RAG implementation using LangChain, Milvus
7
+and vLLM. RAG enhances LLM responses by retrieving relevant context
8
+from a document collection.
9
+
10
+Features:
11
+- Web content loading and chunking
12
+- Vector storage with Milvus
13
+- Embedding generation with vLLM
14
+- Question answering with context
15
+
16
+Prerequisites:
17
+1. Install dependencies:
18
+ pip install -U vllm \
19
+ langchain_milvus langchain_openai \
20
+ langchain_community beautifulsoup4 \
21
+ langchain-text-splitters
22
+
23
+2. Start services:
24
+ # Start embedding service (port 8000)
25
+ vllm serve ssmits/Qwen2-7B-Instruct-embed-base
26
+
27
+ # Start chat service (port 8001)
28
+ vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
29
+
30
+Usage:
31
+ python retrieval_augmented_generation_with_langchain.py
32
+
33
+Notes:
34
+ - Ensure both vLLM services are running before executing
35
+ - Default ports: 8000 (embedding), 8001 (chat)
36
+ - First run may take time to download models
37
+"""
38
+
39
+import argparse
40
+from argparse import Namespace
41
+from typing import Any
42
+
43
+from langchain_community.document_loaders import WebBaseLoader
44
+from langchain_core.documents import Document
45
+from langchain_core.output_parsers import StrOutputParser
46
+from langchain_core.prompts import PromptTemplate
47
+from langchain_core.runnables import RunnablePassthrough
48
+from langchain_milvus import Milvus
49
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
50
+from langchain_text_splitters import RecursiveCharacterTextSplitter
51
+
52
+
53
+def load_and_split_documents(config: dict[str, Any]):
54
+ """
55
+ Load and split documents from web URL
56
+ """
57
+ try:
58
+ loader = WebBaseLoader(web_paths=(config["url"], ))
59
+ docs = loader.load()
60
+
61
+ text_splitter = RecursiveCharacterTextSplitter(
62
+ chunk_size=config["chunk_size"],
63
+ chunk_overlap=config["chunk_overlap"],
64
+ )
65
+ return text_splitter.split_documents(docs)
66
+ except Exception as e:
67
+ print(f"Error loading document from {config['url']}: {str(e)}")
68
+ raise
69
+
70
+
71
+def init_vectorstore(config: dict[str, Any], documents: list[Document]):
72
+ """
73
+ Initialize vector store with documents
74
+ """
75
+ return Milvus.from_documents(
76
+ documents=documents,
77
+ embedding=OpenAIEmbeddings(
78
+ model=config["embedding_model"],
79
+ openai_api_key=config["vllm_api_key"],
80
+ openai_api_base=config["vllm_embedding_endpoint"],
81
+ ),
82
+ connection_args={"uri": config["uri"]},
83
+ drop_old=True,
84
+ )
85
+
86
+
87
+def init_llm(config: dict[str, Any]):
88
+ """
89
+ Initialize llm
90
+ """
91
+ return ChatOpenAI(
92
+ model=config["chat_model"],
93
+ openai_api_key=config["vllm_api_key"],
94
+ openai_api_base=config["vllm_chat_endpoint"],
95
+ )
96
+
97
+
98
+def get_qa_prompt():
99
+ """
100
+ Get question answering prompt template
101
+ """
102
+ template = """You are an assistant for question-answering tasks.
103
+Use the following pieces of retrieved context to answer the question.
104
+If you don't know the answer, just say that you don't know.
105
+Use three sentences maximum and keep the answer concise.
106
+Question: {question}
107
+Context: {context}
108
+Answer:
109
+"""
110
+ return PromptTemplate.from_template(template)
111
+
112
+
113
+def format_docs(docs: list[Document]):
114
+ """
115
+ Format documents for prompt
116
+ """
117
+ return "\n\n".join(doc.page_content for doc in docs)
118
+
119
+
120
+def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
121
+ """
122
+ Set up question answering chain
123
+ """
124
+ return ({
125
+ "context": retriever | format_docs,
126
+ "question": RunnablePassthrough(),
127
+ }
128
+ | prompt
129
+ | llm
130
+ | StrOutputParser())
131
+
132
+
133
+def get_parser() -> argparse.ArgumentParser:
134
+ """
135
+ Parse command line arguments
136
+ """
137
+ parser = argparse.ArgumentParser(description='RAG with vLLM and langchain')
138
+
139
+ # Add command line arguments
140
+ parser.add_argument('--vllm-api-key',
141
+ default="EMPTY",
142
+ help='API key for vLLM compatible services')
143
+ parser.add_argument('--vllm-embedding-endpoint',
144
+ default="http://localhost:8000/v1",
145
+ help='Base URL for embedding service')
146
+ parser.add_argument('--vllm-chat-endpoint',
147
+ default="http://localhost:8001/v1",
148
+ help='Base URL for chat service')
149
+ parser.add_argument('--uri',
150
+ default="./milvus.db",
151
+ help='URI for Milvus database')
152
+ parser.add_argument(
153
+ '--url',
154
+ default=("https://docs.vllm.ai/en/latest/getting_started/"
155
+ "quickstart.html"),
156
+ help='URL of the document to process')
157
+ parser.add_argument('--embedding-model',
158
+ default="ssmits/Qwen2-7B-Instruct-embed-base",
159
+ help='Model name for embeddings')
160
+ parser.add_argument('--chat-model',
161
+ default="qwen/Qwen1.5-0.5B-Chat",
162
+ help='Model name for chat')
163
+ parser.add_argument('-i',
164
+ '--interactive',
165
+ action='store_true',
166
+ help='Enable interactive Q&A mode')
167
+ parser.add_argument('-k',
168
+ '--top-k',
169
+ type=int,
170
+ default=3,
171
+ help='Number of top results to retrieve')
172
+ parser.add_argument('-c',
173
+ '--chunk-size',
174
+ type=int,
175
+ default=1000,
176
+ help='Chunk size for document splitting')
177
+ parser.add_argument('-o',
178
+ '--chunk-overlap',
179
+ type=int,
180
+ default=200,
181
+ help='Chunk overlap for document splitting')
182
+
183
+ return parser
184
+
185
+
186
+def init_config(args: Namespace):
187
+ """
188
+ Initialize configuration settings from command line arguments
189
+ """
190
+
191
+ return {
192
+ "vllm_api_key": args.vllm_api_key,
193
+ "vllm_embedding_endpoint": args.vllm_embedding_endpoint,
194
+ "vllm_chat_endpoint": args.vllm_chat_endpoint,
195
+ "uri": args.uri,
196
+ "embedding_model": args.embedding_model,
197
+ "chat_model": args.chat_model,
198
+ "url": args.url,
199
+ "chunk_size": args.chunk_size,
200
+ "chunk_overlap": args.chunk_overlap,
201
+ "top_k": args.top_k
202
+ }
203
+
204
+
205
+def main():
206
+ # Parse command line arguments
207
+ args = get_parser().parse_args()
208
+
209
+ # Initialize configuration
210
+ config = init_config(args)
211
+
212
+ # Load and split documents
213
+ documents = load_and_split_documents(config)
214
+
215
+ # Initialize vector store and retriever
216
+ vectorstore = init_vectorstore(config, documents)
217
+ retriever = vectorstore.as_retriever(search_kwargs={"k": config["top_k"]})
218
+
219
+ # Initialize llm and prompt
220
+ llm = init_llm(config)
221
+ prompt = get_qa_prompt()
222
+
223
+ # Set up QA chain
224
+ qa_chain = create_qa_chain(retriever, llm, prompt)
225
+
226
+ # Interactive mode
227
+ if args.interactive:
228
+ print("\nWelcome to Interactive Q&A System!")
229
+ print("Enter 'q' or 'quit' to exit.")
230
+
231
+ while True:
232
+ question = input("\nPlease enter your question: ")
233
+ if question.lower() in ['q', 'quit']:
234
+ print("\nThank you for using! Goodbye!")
235
+ break
236
+
237
+ output = qa_chain.invoke(question)
238
+ print(output)
239
+ else:
240
+ # Default single question mode
241
+ question = ("How to install vLLM?")
242
+ output = qa_chain.invoke(question)
243
+ print("-" * 50)
244
+ print(output)
245
+ print("-" * 50)
246
+
247
+
248
+if __name__ == "__main__":
249
+ main()
examples/online_serving/retrieval_augmented_generation_with_llamaindex.py
new
+217
@@ -0,0 +1,217 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
4
+================================================================
5
+
6
+This script demonstrates a RAG system using:
7
+- LlamaIndex: For document indexing and retrieval
8
+- Milvus: As vector store backend
9
+- vLLM: For embedding and text generation
10
+
11
+Features:
12
+1. Document Loading & Processing
13
+2. Embedding & Storage
14
+3. Query Processing
15
+
16
+Requirements:
17
+1. Install dependencies:
18
+pip install llama-index llama-index-readers-web \
19
+ llama-index-llms-openai-like \
20
+ llama-index-embeddings-openai-like \
21
+ llama-index-vector-stores-milvus \
22
+
23
+2. Start services:
24
+ # Start embedding service (port 8000)
25
+ vllm serve ssmits/Qwen2-7B-Instruct-embed-base
26
+
27
+ # Start chat service (port 8001)
28
+ vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
29
+
30
+Usage:
31
+ python retrieval_augmented_generation_with_llamaindex.py
32
+
33
+Notes:
34
+ - Ensure both vLLM services are running before executing
35
+ - Default ports: 8000 (embedding), 8001 (chat)
36
+ - First run may take time to download models
37
+"""
38
+import argparse
39
+from argparse import Namespace
40
+from typing import Any
41
+
42
+from llama_index.core import Settings, StorageContext, VectorStoreIndex
43
+from llama_index.core.node_parser import SentenceSplitter
44
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
45
+from llama_index.llms.openai_like import OpenAILike
46
+from llama_index.readers.web import SimpleWebPageReader
47
+from llama_index.vector_stores.milvus import MilvusVectorStore
48
+
49
+
50
+def init_config(args: Namespace):
51
+ """Initialize configuration with command line arguments"""
52
+ return {
53
+ "url": args.url,
54
+ "embedding_model": args.embedding_model,
55
+ "chat_model": args.chat_model,
56
+ "vllm_api_key": args.vllm_api_key,
57
+ "embedding_endpoint": args.embedding_endpoint,
58
+ "chat_endpoint": args.chat_endpoint,
59
+ "db_path": args.db_path,
60
+ "chunk_size": args.chunk_size,
61
+ "chunk_overlap": args.chunk_overlap,
62
+ "top_k": args.top_k
63
+ }
64
+
65
+
66
+def load_documents(url: str) -> list:
67
+ """Load and process web documents"""
68
+ return SimpleWebPageReader(html_to_text=True).load_data([url])
69
+
70
+
71
+def setup_models(config: dict[str, Any]):
72
+ """Configure embedding and chat models"""
73
+ Settings.embed_model = OpenAILikeEmbedding(
74
+ api_base=config["embedding_endpoint"],
75
+ api_key=config["vllm_api_key"],
76
+ model_name=config["embedding_model"],
77
+ )
78
+
79
+ Settings.llm = OpenAILike(
80
+ model=config["chat_model"],
81
+ api_key=config["vllm_api_key"],
82
+ api_base=config["chat_endpoint"],
83
+ context_window=128000,
84
+ is_chat_model=True,
85
+ is_function_calling_model=False,
86
+ )
87
+
88
+ Settings.transformations = [
89
+ SentenceSplitter(
90
+ chunk_size=config["chunk_size"],
91
+ chunk_overlap=config["chunk_overlap"],
92
+ )
93
+ ]
94
+
95
+
96
+def setup_vector_store(db_path: str) -> MilvusVectorStore:
97
+ """Initialize vector store"""
98
+ sample_emb = Settings.embed_model.get_text_embedding("test")
99
+ print(f"Embedding dimension: {len(sample_emb)}")
100
+ return MilvusVectorStore(uri=db_path, dim=len(sample_emb), overwrite=True)
101
+
102
+
103
+def create_index(documents: list, vector_store: MilvusVectorStore):
104
+ """Create document index"""
105
+ storage_context = StorageContext.from_defaults(vector_store=vector_store)
106
+ return VectorStoreIndex.from_documents(
107
+ documents,
108
+ storage_context=storage_context,
109
+ )
110
+
111
+
112
+def query_document(index: VectorStoreIndex, question: str, top_k: int):
113
+ """Query document with given question"""
114
+ query_engine = index.as_query_engine(similarity_top_k=top_k)
115
+ return query_engine.query(question)
116
+
117
+
118
+def get_parser() -> argparse.ArgumentParser:
119
+ """Parse command line arguments"""
120
+ parser = argparse.ArgumentParser(
121
+ description='RAG with vLLM and LlamaIndex')
122
+
123
+ # Add command line arguments
124
+ parser.add_argument(
125
+ '--url',
126
+ default=("https://docs.vllm.ai/en/latest/getting_started/"
127
+ "quickstart.html"),
128
+ help='URL of the document to process')
129
+ parser.add_argument('--embedding-model',
130
+ default="ssmits/Qwen2-7B-Instruct-embed-base",
131
+ help='Model name for embeddings')
132
+ parser.add_argument('--chat-model',
133
+ default="qwen/Qwen1.5-0.5B-Chat",
134
+ help='Model name for chat')
135
+ parser.add_argument('--vllm-api-key',
136
+ default="EMPTY",
137
+ help='API key for vLLM compatible services')
138
+ parser.add_argument('--embedding-endpoint',
139
+ default="http://localhost:8000/v1",
140
+ help='Base URL for embedding service')
141
+ parser.add_argument('--chat-endpoint',
142
+ default="http://localhost:8001/v1",
143
+ help='Base URL for chat service')
144
+ parser.add_argument('--db-path',
145
+ default="./milvus_demo.db",
146
+ help='Path to Milvus database')
147
+ parser.add_argument('-i',
148
+ '--interactive',
149
+ action='store_true',
150
+ help='Enable interactive Q&A mode')
151
+ parser.add_argument('-c',
152
+ '--chunk-size',
153
+ type=int,
154
+ default=1000,
155
+ help='Chunk size for document splitting')
156
+ parser.add_argument('-o',
157
+ '--chunk-overlap',
158
+ type=int,
159
+ default=200,
160
+ help='Chunk overlap for document splitting')
161
+ parser.add_argument('-k',
162
+ '--top-k',
163
+ type=int,
164
+ default=3,
165
+ help='Number of top results to retrieve')
166
+
167
+ return parser
168
+
169
+
170
+def main():
171
+ # Parse command line arguments
172
+ args = get_parser().parse_args()
173
+
174
+ # Initialize configuration
175
+ config = init_config(args)
176
+
177
+ # Load documents
178
+ documents = load_documents(config["url"])
179
+
180
+ # Setup models
181
+ setup_models(config)
182
+
183
+ # Setup vector store
184
+ vector_store = setup_vector_store(config["db_path"])
185
+
186
+ # Create index
187
+ index = create_index(documents, vector_store)
188
+
189
+ if args.interactive:
190
+ print("\nEntering interactive mode. Type 'quit' to exit.")
191
+ while True:
192
+ # Get user question
193
+ question = input("\nEnter your question: ")
194
+
195
+ # Check for exit command
196
+ if question.lower() in ['quit', 'exit', 'q']:
197
+ print("Exiting interactive mode...")
198
+ break
199
+
200
+ # Get and print response
201
+ print("\n" + "-" * 50)
202
+ print("Response:\n")
203
+ response = query_document(index, question, config["top_k"])
204
+ print(response)
205
+ print("-" * 50)
206
+ else:
207
+ # Single query mode
208
+ question = "How to install vLLM?"
209
+ response = query_document(index, question, config["top_k"])
210
+ print("-" * 50)
211
+ print("Response:\n")
212
+ print(response)
213
+ print("-" * 50)
214
+
215
+
216
+if __name__ == "__main__":
217
+ main()