File size: 4,338 Bytes
4416e3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import os
import json
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_classic.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document
load_dotenv(override=True)
# Configuration
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if GOOGLE_API_KEY:
print(f"RAG Model: API Key loaded (starting with {GOOGLE_API_KEY[:5]}...)")
else:
print("RAG Model: No API Key found in environment!")
DATA_FILE = "preprocessed_docs.json"
PERSIST_DIRECTORY = "chroma_db"
def build_rag_system():
if not os.path.exists(DATA_FILE):
print("Data file not found. Please run preprocess.py first.")
return None
with open(DATA_FILE, "r") as f:
docs_data = json.load(f)
print(f"Loading {len(docs_data)} documents...")
# Prepare documents for LangChain
documents = [Document(page_content=d["content"], metadata=d["metadata"]) for d in docs_data]
# Split documents into chunks (optional, but good practice)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
chunks = text_splitter.split_documents(documents)
print(f"Created {len(chunks)} chunks.")
# Initialize Embeddings
print("Initializing embeddings (HuggingFace)...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
# Initialize Vector Store
print("Building vector store...")
vectorstore = Chroma.from_documents(
documents=chunks,
embedding=embeddings,
persist_directory=PERSIST_DIRECTORY
)
vectorstore.persist()
print("Vector store built and persisted.")
return vectorstore
def get_qa_chain(vectorstore):
if not GOOGLE_API_KEY:
print("Warning: GOOGLE_API_KEY not found. LLM functionality will not work.")
return None
llm = ChatGoogleGenerativeAI(model="gemini-flash-latest", google_api_key=GOOGLE_API_KEY)
# Custom Prompt
template = """You are a helpful and expert Retail Product Assistant.
Context (Product Details & Reviews):
{context}
Rules:
1. If the user says "hi", "hello" or greets you, greet them back warmly and mention 2-3 popular products from the context to get started.
2. Use the provided context to answer specific questions.
3. If the answer is not in the context, politely say you don't have that specific information.
4. Maintain a professional yet friendly tone.
5. Always use Markdown for formatting (bolding, lists, etc.) to make it easy to read.
6. Use bullet points if listing features or pros/cons.
7. IMPORTANT: Convert all prices from USD to Indian Rupees (INR) using an approximate exchange rate of 1 USD = ₹83 INR. Always display the price in INR (₹).
Question: {question}
Answer:"""
QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vectorstore.as_retriever(),
return_source_documents=True,
chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
)
return qa_chain
if __name__ == "__main__":
# Check if vector store exists
if os.path.exists(PERSIST_DIRECTORY):
print("Vector store already exists. Loading...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)
else:
vectorstore = build_rag_system()
if vectorstore:
qa_chain = get_qa_chain(vectorstore)
if qa_chain:
query = "Which Kindle model has the best resolution according to reviews?"
print(f"\nQuestion: {query}")
result = qa_chain({"query": query})
print(f"\nAnswer: {result['result']}")
print("\nSources:")
for doc in result["source_documents"][:2]:
print(f"- {doc.metadata['name']} (Rating: {doc.metadata['rating']})")
|