Spaces:
Runtime error
Runtime error
Download rag_chain.py from StudyOcean/Study-Ocean-AI_Mentor: direct link, hf CLI and curl.
- Browser
- Download file 5.34 kB
-
https://huggingface.co/spaces/StudyOcean/Study-Ocean-AI_Mentor/resolve/main/rag_chain.py
- Command line
-
hf download hf://spaces/StudyOcean/Study-Ocean-AI_Mentor/rag_chain.py
-
curl -L -o rag_chain.py https://huggingface.co/spaces/StudyOcean/Study-Ocean-AI_Mentor/resolve/main/rag_chain.py
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 | |