Spaces:
Sleeping
Sleeping
Download rag_engine.py from adeena43/python-rag-assistant: direct link, hf CLI and curl.
- Browser
- Download file 5.61 kB
-
https://huggingface.co/spaces/adeena43/python-rag-assistant/resolve/main/rag_engine.py
- Command line
-
hf download hf://spaces/adeena43/python-rag-assistant/rag_engine.py
-
curl -L -o rag_engine.py https://huggingface.co/spaces/adeena43/python-rag-assistant/resolve/main/rag_engine.py
5.61 kB
| # rag_engine.py | |
| import os | |
| import warnings | |
| from typing import List | |
| from dotenv import load_dotenv | |
| # Suppress warnings | |
| warnings.filterwarnings('ignore') | |
| # Set environment variables before imports | |
| os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" | |
| os.environ["USER_AGENT"] = "PythonicAI/1.0 (+adinafraz01@gmail.com)" | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| # Load env variables | |
| load_dotenv() | |
| openai_key = os.getenv("OPENAI_API_KEY") | |
| if not openai_key: | |
| raise ValueError("OPENAI_API_KEY not found in environment variables. Please create a .env file with your API key.") | |
| os.environ["OPENAI_API_KEY"] = openai_key | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_openai import OpenAIEmbeddings, ChatOpenAI | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_community.document_loaders import WebBaseLoader | |
| from langchain_core.prompts import PromptTemplate | |
| from langchain.schema import Document | |
| # ------------------------------------------ | |
| # DOCUMENT LOADING (URL ONLY) | |
| # ------------------------------------------ | |
| def load_user_document(url: str) -> List[Document]: | |
| """Load URL into LangChain documents.""" | |
| try: | |
| if not url: | |
| raise ValueError("No URL provided") | |
| # Clean URL | |
| if not url.startswith(('http://', 'https://')): | |
| url = f'https://{url}' | |
| # Use WebBaseLoader with custom headers | |
| headers = { | |
| 'User-Agent': os.environ.get("USER_AGENT", "PythonicAI/1.0") | |
| } | |
| loader = WebBaseLoader( | |
| web_paths=[url], | |
| requests_per_second=1, | |
| requests_kwargs={'headers': headers} | |
| ) | |
| docs = loader.load() | |
| if not docs: | |
| raise ValueError("No content found in the URL") | |
| return docs | |
| except Exception as e: | |
| raise Exception(f"Error loading URL: {str(e)}") | |
| # ------------------------------------------ | |
| # BUILD RETRIEVER FROM DOCUMENTS | |
| # ------------------------------------------ | |
| def build_retriever_from_docs(docs: List[Document]): | |
| """Build FAISS vector store from documents.""" | |
| if not docs: | |
| raise ValueError("No documents provided") | |
| # Use the same splitter settings as your example | |
| splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=1000, | |
| chunk_overlap=200 | |
| ) | |
| chunks = splitter.split_documents(docs) | |
| # Use the same embedding model as your example | |
| embeddings = OpenAIEmbeddings(model="text-embedding-3-small") | |
| # Build vector store | |
| vector_store = FAISS.from_documents(chunks, embeddings) | |
| return vector_store.as_retriever( | |
| search_type="similarity", | |
| search_kwargs={"k": 6} | |
| ) | |
| # ------------------------------------------ | |
| # PYTHON DOMAIN CHECK | |
| # ------------------------------------------ | |
| _llm = None | |
| def get_llm(): | |
| """Singleton pattern for LLM to avoid re-initialization.""" | |
| global _llm | |
| if _llm is None: | |
| _llm = ChatOpenAI( | |
| model="gpt-4o-mini", | |
| temperature=0.2, | |
| max_tokens=1000, | |
| request_timeout=60 | |
| ) | |
| return _llm | |
| def is_python_related(question: str) -> bool: | |
| """Check if question is Python-related.""" | |
| python_keywords = [ | |
| 'python', 'numpy', 'pandas', 'matplotlib', 'seaborn', | |
| 'machine learning', 'deep learning', 'tensorflow', 'pytorch', | |
| 'scikit-learn', 'sklearn', 'dataframe', 'array', 'plot', | |
| 'import', 'def ', 'class ', 'function', 'method', 'library', | |
| 'package', 'module', 'pip install', 'conda install', 'jupyter', | |
| 'data', 'science', 'analysis', 'visualization', 'programming', | |
| 'code', 'script', 'algorithm', 'model', 'training', 'prediction' | |
| ] | |
| # Quick keyword check first | |
| question_lower = question.lower() | |
| for keyword in python_keywords: | |
| if keyword in question_lower: | |
| return True | |
| return False | |
| # ------------------------------------------ | |
| # ANSWER QUESTION | |
| # ------------------------------------------ | |
| # Define the prompt template (same as your example) | |
| prompt = PromptTemplate( | |
| template=""" | |
| You are a helpful python assistant. | |
| Answer ONLY from the provided text context. | |
| Display answer in a beautiful and user understandable format with proper formatting. | |
| Use bullet points, numbered lists, or code blocks when appropriate. | |
| If the context is insufficient, just say you don't know. | |
| Context: {context} | |
| Question: {question} | |
| """, | |
| input_variables=['context', 'question'] | |
| ) | |
| def answer_question(question: str, retriever): | |
| """Answer question using RAG pipeline.""" | |
| # Validate topic | |
| if not is_python_related(question): | |
| return "I can only answer Python and data science related questions based on the provided documentation." | |
| try: | |
| # Retrieve relevant documents | |
| docs = retriever.invoke(question) | |
| if not docs: | |
| return "I couldn't find relevant information in the documentation to answer this question." | |
| # Combine context from retrieved documents | |
| context = "\n\n".join(doc.page_content for doc in docs) | |
| # Create final prompt | |
| final_prompt = prompt.invoke({ | |
| "context": context, | |
| "question": question | |
| }) | |
| # Get answer from LLM | |
| llm = get_llm() | |
| answer = llm.invoke(final_prompt) | |
| return answer.content | |
| except Exception as e: | |
| return f"Error generating answer: {str(e)}" |