Spaces:
Sleeping
Sleeping
File size: 4,403 Bytes
a4538e5 1b9900c a4538e5 1b9900c a4538e5 | 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 | import os
from dataclasses import dataclass
from typing import List, Optional
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage
from src.vectorstore import FaissVectorStore
load_dotenv()
@dataclass
class RetrievalResult:
index: int
distance: float
text: Optional[str]
class RAGSearch:
def __init__(
self,
persist_dir: str = "faiss_store",
embedding_model: str = "all-MiniLM-L6-v2",
llm_model: str = "llama-3.1-8b-instant",
):
# Vector store setup
self.vectorstore = FaissVectorStore(persist_dir=persist_dir, embedding_model=embedding_model)
faiss_path = os.path.join(persist_dir, "faiss.index")
meta_path = os.path.join(persist_dir, "metadata.pkl")
if not (os.path.exists(faiss_path) and os.path.exists(meta_path)):
# Build from local 'Research/data' directory if index doesn't exist
from src.data_loader import load_all_documents
# Try multiple possible data directories
data_dirs = ["Research/data", "data", "Data"]
docs = []
for data_dir in data_dirs:
if os.path.exists(data_dir):
print(f"[INFO] Checking for documents in: {data_dir}")
docs = load_all_documents(data_dir)
if docs:
print(f"[INFO] Found {len(docs)} documents in {data_dir}")
break
if not docs:
print("[WARNING] No documents found in any data directory. Vector store will be empty.")
print("[INFO] Please add documents to 'Research/data/', 'data/', or 'Data/' directory.")
# Create empty index for now
self.vectorstore.index = None
self.vectorstore.metadata = []
else:
self.vectorstore.build_from_documents(docs)
else:
self.vectorstore.load()
# LLM setup
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
raise ValueError("GROQ_API_KEY missing in environment")
self.llm_model = llm_model
self.embedding_model = embedding_model
self.llm = ChatGroq(api_key=groq_api_key, model=llm_model, temperature=0.1)
print(f"[INFO] Groq LLM initialized: {llm_model}")
def retrieve(self, query: str, top_k: int = 5) -> List[RetrievalResult]:
# Check if vector store is empty
if self.vectorstore.index is None or len(self.vectorstore.metadata) == 0:
print("[WARNING] Vector store is empty. No documents to search.")
return []
results = self.vectorstore.query(query_text=query, top_k=top_k)
out: List[RetrievalResult] = []
for r in results:
text = r["metadata"]["texts"] if r.get("metadata") and r["metadata"].get("texts") else None
out.append(RetrievalResult(index=int(r["index"]), distance=float(r["distance"]), text=text))
return out
def summarize(self, query: str, retrieved: List[RetrievalResult]) -> str:
texts = [r.text for r in retrieved if r.text]
context = "\n\n".join(texts)
# Using proper message formatting for better LLM interaction
system_message = SystemMessage(content="You are a helpful assistant. Use the provided context to answer the user's question if the information is present. If the answer is not in the context, or if the context is empty, answer the question using your own knowledge.")
human_message = HumanMessage(content=f"""
Context:
{context}
Query: {query}
Answer:
""")
try:
response = self.llm.invoke([system_message, human_message])
return response.content
except Exception as e:
return f"Error generating response: {str(e)}. Please check your GROQ_API_KEY is set correctly."
def search_and_summarize(self, query: str, top_k: int = 5) -> str:
retrieved = self.retrieve(query, top_k=top_k)
return self.summarize(query, retrieved)
if __name__ == "__main__":
rag_search = RAGSearch()
query = "What is Database Management System?"
summary = rag_search.search_and_summarize(query, top_k=3)
print("Summary:", summary) |