| import os
|
| import pandas as pd
|
| import json
|
| from langchain_community.vectorstores import FAISS
|
| from langchain_huggingface import HuggingFaceEmbeddings
|
| from langchain_text_splitters import RecursiveCharacterTextSplitter
|
| from langchain_core.documents import Document
|
| from transformers import pipeline
|
|
|
| class RAGSystem:
|
| def __init__(self, data_dir):
|
| self.data_dir = data_dir
|
| self.vector_store = None
|
| self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
|
| self.llm_pipeline = None
|
|
|
| def load_documents(self):
|
| docs = []
|
|
|
|
|
| reviews_path = os.path.join(self.data_dir, 'customer_reviews.csv')
|
| if os.path.exists(reviews_path):
|
| df_reviews = pd.read_csv(reviews_path)
|
| for _, row in df_reviews.iterrows():
|
| content = f"Product: {row.get('Product', 'Unknown')}\nDate: {row.get('Date', '')}\nRating: {row.get('Rating', '')}\nReview: {row.get('ReviewText', '')}"
|
| metadata = {"source": "customer_reviews", "product": row.get('Product', 'Unknown')}
|
| docs.append(Document(page_content=content, metadata=metadata))
|
|
|
|
|
| logs_path = os.path.join(self.data_dir, 'web_logs.json')
|
| if os.path.exists(logs_path):
|
| with open(logs_path, 'r') as f:
|
| logs_data = json.load(f)
|
| for log in logs_data:
|
| content = f"Log Timestamp: {log.get('timestamp', '')}\nAction: {log.get('action', '')}\nPage: {log.get('page', '')}\nUser: {log.get('user_id', '')}"
|
| metadata = {"source": "web_logs"}
|
| docs.append(Document(page_content=content, metadata=metadata))
|
|
|
|
|
| sales_path = os.path.join(self.data_dir, 'sales_data.csv')
|
| if os.path.exists(sales_path):
|
| df_sales = pd.read_csv(sales_path)
|
|
|
| summary = df_sales.groupby(['Product', 'Region'])['TotalPrice'].sum().reset_index()
|
| for _, row in summary.iterrows():
|
| content = f"Sales Summary:\nProduct: {row['Product']}\nRegion: {row['Region']}\nTotal Revenue: ${row['TotalPrice']:.2f}"
|
| metadata = {"source": "sales_summary"}
|
| docs.append(Document(page_content=content, metadata=metadata))
|
|
|
| return docs
|
|
|
| def build_index(self):
|
| docs = self.load_documents()
|
| text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
| splits = text_splitter.split_documents(docs)
|
|
|
| if splits:
|
| self.vector_store = FAISS.from_documents(splits, self.embeddings)
|
| return True
|
| return False
|
|
|
| def init_llm(self):
|
|
|
|
|
| try:
|
| self.llm_pipeline = pipeline("text2text-generation", model="google/flan-t5-small")
|
| except Exception as e:
|
| print(f"Error loading LLM: {e}")
|
| self.llm_pipeline = None
|
|
|
| def query(self, user_query, k=3):
|
| if not self.vector_store:
|
| return {
|
| "answer": "System not initialized. Please build the index first.",
|
| "context": []
|
| }
|
|
|
|
|
| docs = self.vector_store.similarity_search(user_query, k=k)
|
| context_text = "\n\n".join([d.page_content for d in docs])
|
|
|
|
|
| if self.llm_pipeline:
|
| prompt = f"Summarize the following context to answer the question. \n\nContext:\n{context_text}\n\nQuestion: {user_query}\n\nAnswer:"
|
|
|
| if len(prompt) > 2048:
|
| prompt = prompt[:2048]
|
|
|
| result = self.llm_pipeline(prompt, max_length=200, do_sample=False)
|
| answer = result[0]['generated_text']
|
| else:
|
| answer = "LLM not loaded. Displaying retrieved context only."
|
|
|
| return {
|
| "answer": answer,
|
| "context": docs
|
| }
|
|
|
