@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 240 lines 7.21 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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
40 import argparse
41 from argparse import Namespace
42 from typing import Any
43
44 from langchain_community.document_loaders import WebBaseLoader
45 from langchain_core.documents import Document
46 from langchain_core.output_parsers import StrOutputParser
47 from langchain_core.prompts import PromptTemplate
48 from langchain_core.runnables import RunnablePassthrough
49 from langchain_milvus import Milvus
50 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
51 from langchain_text_splitters import RecursiveCharacterTextSplitter
52
53
54 def load_and_split_documents(config: dict[str, Any]):
55 """Load and split documents from web URL."""
56 try:
57 loader = WebBaseLoader(web_paths=(config["url"],))
58 docs = loader.load()
59
60 text_splitter = RecursiveCharacterTextSplitter(
61 chunk_size=config["chunk_size"],
62 chunk_overlap=config["chunk_overlap"],
63 )
64 return text_splitter.split_documents(docs)
65 except Exception as e:
66 print(f"Error loading document from {config['url']}: {str(e)}")
67 raise
68
69
70 def init_vectorstore(config: dict[str, Any], documents: list[Document]):
71 """Initialize vector store with documents."""
72 return Milvus.from_documents(
73 documents=documents,
74 embedding=OpenAIEmbeddings(
75 model=config["embedding_model"],
76 openai_api_key=config["vllm_api_key"],
77 openai_api_base=config["vllm_embedding_endpoint"],
78 ),
79 connection_args={"uri": config["uri"]},
80 drop_old=True,
81 )
82
83
84 def init_llm(config: dict[str, Any]):
85 """Initialize llm."""
86 return ChatOpenAI(
87 model=config["chat_model"],
88 openai_api_key=config["vllm_api_key"],
89 openai_api_base=config["vllm_chat_endpoint"],
90 )
91
92
93 def get_qa_prompt():
94 """Get question answering prompt template."""
95 template = """You are an assistant for question-answering tasks.
96 Use the following pieces of retrieved context to answer the question.
97 If you don't know the answer, just say that you don't know.
98 Use three sentences maximum and keep the answer concise.
99 Question: {question}
100 Context: {context}
101 Answer:
102 """
103 return PromptTemplate.from_template(template)
104
105
106 def format_docs(docs: list[Document]):
107 """Format documents for prompt."""
108 return "\n\n".join(doc.page_content for doc in docs)
109
110
111 def create_qa_chain(retriever: Any, llm: ChatOpenAI, prompt: PromptTemplate):
112 """Set up question answering chain."""
113 return (
114 {
115 "context": retriever | format_docs,
116 "question": RunnablePassthrough(),
117 }
118 | prompt
119 | llm
120 | StrOutputParser()
121 )
122
123
124 def get_parser() -> argparse.ArgumentParser:
125 """Parse command line arguments."""
126 parser = argparse.ArgumentParser(description="RAG with vLLM and langchain")
127
128 # Add command line arguments
129 parser.add_argument(
130 "--vllm-api-key", default="EMPTY", help="API key for vLLM compatible services"
131 )
132 parser.add_argument(
133 "--vllm-embedding-endpoint",
134 default="http://localhost:8000/v1",
135 help="Base URL for embedding service",
136 )
137 parser.add_argument(
138 "--vllm-chat-endpoint",
139 default="http://localhost:8001/v1",
140 help="Base URL for chat service",
141 )
142 parser.add_argument("--uri", default="./milvus.db", help="URI for Milvus database")
143 parser.add_argument(
144 "--url",
145 default=("https://docs.vllm.ai/en/latest/getting_started/quickstart.html"),
146 help="URL of the document to process",
147 )
148 parser.add_argument(
149 "--embedding-model",
150 default="ssmits/Qwen2-7B-Instruct-embed-base",
151 help="Model name for embeddings",
152 )
153 parser.add_argument(
154 "--chat-model", default="qwen/Qwen1.5-0.5B-Chat", help="Model name for chat"
155 )
156 parser.add_argument(
157 "-i", "--interactive", action="store_true", help="Enable interactive Q&A mode"
158 )
159 parser.add_argument(
160 "-k", "--top-k", type=int, default=3, help="Number of top results to retrieve"
161 )
162 parser.add_argument(
163 "-c",
164 "--chunk-size",
165 type=int,
166 default=1000,
167 help="Chunk size for document splitting",
168 )
169 parser.add_argument(
170 "-o",
171 "--chunk-overlap",
172 type=int,
173 default=200,
174 help="Chunk overlap for document splitting",
175 )
176
177 return parser
178
179
180 def init_config(args: Namespace):
181 """Initialize configuration settings from command line arguments."""
182 return {
183 "vllm_api_key": args.vllm_api_key,
184 "vllm_embedding_endpoint": args.vllm_embedding_endpoint,
185 "vllm_chat_endpoint": args.vllm_chat_endpoint,
186 "uri": args.uri,
187 "embedding_model": args.embedding_model,
188 "chat_model": args.chat_model,
189 "url": args.url,
190 "chunk_size": args.chunk_size,
191 "chunk_overlap": args.chunk_overlap,
192 "top_k": args.top_k,
193 }
194
195
196 def main():
197 # Parse command line arguments
198 args = get_parser().parse_args()
199
200 # Initialize configuration
201 config = init_config(args)
202
203 # Load and split documents
204 documents = load_and_split_documents(config)
205
206 # Initialize vector store and retriever
207 vectorstore = init_vectorstore(config, documents)
208 retriever = vectorstore.as_retriever(search_kwargs={"k": config["top_k"]})
209
210 # Initialize llm and prompt
211 llm = init_llm(config)
212 prompt = get_qa_prompt()
213
214 # Set up QA chain
215 qa_chain = create_qa_chain(retriever, llm, prompt)
216
217 # Interactive mode
218 if args.interactive:
219 print("\nWelcome to Interactive Q&A System!")
220 print("Enter 'q' or 'quit' to exit.")
221
222 while True:
223 question = input("\nPlease enter your question: ")
224 if question.lower() in ["q", "quit"]:
225 print("\nThank you for using! Goodbye!")
226 break
227
228 output = qa_chain.invoke(question)
229 print(output)
230 else:
231 # Default single question mode
232 question = "How to install vLLM?"
233 output = qa_chain.invoke(question)
234 print("-" * 50)
235 print(output)
236 print("-" * 50)
237
238
239 if __name__ == "__main__":
240 main()