Study-Ocean-AI_Mentor / rag_chain.py
Mohit Raj
Update rag_chain.py
575816f verified
Raw History Blame Contribute Delete
5.34 kB
"""
RAG Chain module for Study Ocean Student Mentor.
Optimized for Groq (Llama 3.3) and High-Speed Streaming.
"""
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import json
import config
# ─── System Prompt for Student Mentor ────────────────────────
STUDENT_MENTOR_PROMPT = PromptTemplate(
input_variables=["context", "question"],
template="""You are **Study Ocean AI Mentor** β€” a friendly, patient, and knowledgeable tutor who helps students understand their study material deeply.
Your role:
- You solve student doubts with clear, step-by-step explanations
- You use simple language that students can easily understand
- You provide examples and analogies when helpful
- You reference the study material when answering
- If the answer is in the provided context, use it; if not, say so honestly
- You encourage students and make learning enjoyable
**Important Guidelines:**
1. Start with a brief, direct answer, then explain in detail
2. Use bullet points, numbered steps, or tables for clarity
3. If it's a math/science problem, show the solution step by step
4. If it's a theory question, break it into key points
5. Mention which topic/chapter the information relates to when possible
6. 🌟 Use relevant emojis generously in every response!
---
πŸ“š **Relevant Study Material:**
{context}
---
πŸ™‹ **Student's Question:**
{question}
---
πŸŽ“ **Mentor's Response:**
"""
)
def load_vector_store() -> FAISS:
"""Load the pre-built FAISS vector store from disk. Fallback to ingestion if it fails."""
embeddings = HuggingFaceEmbeddings(
model_name=config.EMBEDDING_MODEL,
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True},
)
try:
print("πŸ“š Loading vector store...")
vector_store = FAISS.load_local(
config.VECTOR_STORE_PATH,
embeddings,
allow_dangerous_deserialization=True,
)
print(" βœ… Vector store loaded!")
except Exception as e:
print(f" ⚠ Failed to load vector store: {e}")
print(" πŸ”„ Automatically rebuilding index from PDFs...")
from ingest import run_ingestion
run_ingestion()
vector_store = FAISS.load_local(
config.VECTOR_STORE_PATH,
embeddings,
allow_dangerous_deserialization=True,
)
return vector_store
def build_rag_chain(vector_store: FAISS) -> RetrievalQA:
"""Build the RAG chain with Groq LLM."""
retriever = vector_store.as_retriever(
search_type="similarity",
search_kwargs={"k": config.TOP_K_RESULTS},
)
# Initialize Groq LLM (Super Fast)
from langchain_groq import ChatGroq
import re
api_key = re.sub(r'[^a-zA-Z0-9_\-]', '', str(config.GROQ_API_KEY))
print(f" ⚑ Groq Engine active. Key length: {len(api_key)}")
llm = ChatGroq(
model_name=config.LLM_MODEL,
groq_api_key=api_key,
temperature=config.TEMPERATURE,
max_tokens=config.MAX_OUTPUT_TOKENS,
)
chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=retriever,
return_source_documents=True,
chain_type_kwargs={"prompt": STUDENT_MENTOR_PROMPT},
)
return chain
def ask_question(chain: RetrievalQA, question: str) -> dict:
result = chain.invoke({"query": question})
sources = []
seen = set()
for doc in result.get("source_documents", []):
source_name = doc.metadata.get("source", "Unknown")
subject = doc.metadata.get("subject", "General")
key = f"{subject}/{source_name}"
if key not in seen:
seen.add(key)
sources.append({"file": source_name, "subject": subject})
return {"answer": result["result"], "sources": sources}
def stream_question(chain: RetrievalQA, question: str):
"""Generator that streams responses from Groq."""
try:
print(f"πŸ“‘ Request: '{question[:50]}...'")
docs = chain.retriever.get_relevant_documents(question)
context = "\n\n".join([doc.page_content for doc in docs])
sources = []
seen = set()
for doc in docs:
name = doc.metadata.get("source", "Unknown")
sub = doc.metadata.get("subject", "General")
if f"{sub}/{name}" not in seen:
seen.add(f"{sub}/{name}")
sources.append({"file": name, "subject": sub})
llm = chain.combine_documents_chain.llm_chain.llm
prompt = STUDENT_MENTOR_PROMPT.format(context=context, question=question)
yield f"DATA:{json.dumps({'sources': sources})}\n\n"
for chunk in llm.stream(prompt):
yield chunk.content
print(" βœ… Stream complete.")
except Exception as e:
print(f" ❌ Error: {str(e)}")
yield f"\n\n**Error:** {str(e)}"
# Singleton
_chain_instance = None
def get_chain():
global _chain_instance
if _chain_instance is None:
vector_store = load_vector_store()
_chain_instance = build_rag_chain(vector_store)
return _chain_instance